Kaynağa Gözat

Added linux-like port-selection-strategy

Denis Waßmann 6 yıl önce
ebeveyn
işleme
a7275fb075
1 değiştirilmiş dosya ile 86 ekleme ve 5 silme
  1. 86 5
      code/ID2TLib/Ports.py

+ 86 - 5
code/ID2TLib/Ports.py

@@ -24,11 +24,11 @@ class PortSelectionStrategy:
 		
 		# that function will always return a one higher counter than before,
 		# restarting from the start once it reached the highest value
-		def __call__(port_range):
+		def __call__(self, port_range, *args):
 			if self.counter == -1:
 				self.counter = port_range.start
 			
-			port = counter
+			port = self.counter
 			
 			self.counter += 1
 			if self.counter == port_range.stop:
@@ -36,9 +36,90 @@ class PortSelectionStrategy:
 			
 			return port
 	class random:
-		def __call__(port_range):
+		def __call__(port_range, *args):
 			return random.randrange(port_range.start, port_range.stop)
 
+	class linux_kernel:
+		"""
+		A port-selectioin-strategy oriented on the linux-kernel
+		The implementation follows https://github.com/torvalds/linux/blob/master/net/ipv4/inet_connection_sock.c#L173
+		as much as possible when converting from one language to another (The newest file was used
+		by the time of writing, make sure you select the correct one when following the link!)
+		"""
+
+		def __call__(self, port_range: range, port_selector, *args):
+			"""
+			This method is an attempt to map a c-function to python. To solve the goto-problem
+			while-true's have been added. Both of the while-true's are placed where the original
+			had a label to jump to. break's and continue's are set to preserve the original
+			control flow. Another method could have been used to rewrite the c-code, however this
+			was chosen to preserve the similarity between this and the original
+
+			:param port_range: the port range to choose from
+			:param port_selector: the port selector that tells which ports are in use
+			:param args: Not used for now
+			:return: A port number
+			"""
+			port = 0
+			low, high = port_range.start, port_range.stop
+
+			# this var tells us if we should use the upper or lower port-range-half, or the whole range if
+			# this var is None. The original was an enum of the values 0, 1 and 2. But I think an Optional[bool]
+			# is more clear
+			# None: use whole range, True: use lower half, False: use upper half
+			attempt_half = True
+
+			high += 1  # line 186 in the original file
+			while True:
+				if high - low < 4:
+					attempt_half = None
+				if attempt_half is not None:
+					# appearently a fast method to find a number close to the real half
+					# unless the difference between high and low is 4 (see above, note the 2-shift below)
+					# this does not work
+					half = low + (((high - low) >> 2) << 1)
+
+					if attempt_half:
+						high = half
+					else:
+						low = half
+
+				remaining = high - low
+				if remaining > 1:
+					remaining &= ~1 # flip the 1-bit
+
+				offset = random.randrange(0, remaining)
+				offset |= 1;
+
+				attempt_half_before = attempt_half # slight hack to keep track of change
+				while True:
+					port = low + offset
+
+					for i in range(0, remaining, 2):
+						if port >= high:
+							port -= remaining
+
+						if port_selector.is_port_in_use(port):
+							port += 2
+							continue
+
+						return port
+
+					offset -= 1
+					if not (offset & 1):
+						continue
+
+					if attempt_half:
+						attempt_half = False
+						break
+
+				if attempt_half_before: # we still got ports to search, attemp_half was just set to False
+					continue
+				if not attempt_half: # the port-range is exhausted
+					break
+
+			raise ValueError("Could not find suitable port")
+
 class PortSelector:
 	"""
 	This class simulates a port-selection-process. Instances keep a list of port-numbers they generated so
@@ -55,7 +136,7 @@ class PortSelector:
 		
 		if len(port_range) == 0:
 			raise ValueError("cannot choose from an empty range")
-		if x.start not in range(1, 65536) or x.stop not in range(1, 65536):
+		if port_range.start not in range(1, 65536) or port_range.stop not in range(1, 65536 + 1):
 			raise ValueError("port_range is no subset of the valid port-range")
 		
 		self.port_range = port_range
@@ -70,7 +151,7 @@ class PortSelector:
 			raise RuntimeError("All %i port numbers were already generated, no more can be generated" % len(self.port_range))
 		
 		while True:
-			port = self._select_port(self.port_range)
+			port = self._select_port(self.port_range, self)
 			
 			if port not in self.generated:
 				self.generated.append(port)