implement naive port prediction

port_prediction.nim

@@ -1,6 +1,165 @@
-from net import Port
+import algorithm
+import net
+import sequtils
+import unittest
 
-proc predictPortRange*(dstPort: Port, probedDstPorts: seq[Port]): seq[Port] =
-  # TODO: do real port prediction
-  result = @[dstPort]
+const RandomPortCount = 1000
 
+proc min(a, b: uint16): uint16 =
+  min(a.int32, b.int32).uint16
+
+proc toUint16(p: Port): uint16 = uint16(p)
+
+proc toPort(u: uint16): Port = Port(u)
+
+proc addOffset(port: uint16, offset: uint16, minValue = 1024'u16,
+               maxValue = uint16.high): uint16 =
+  assert(port >= minValue)
+  assert(port <= maxValue)
+  let distanceToMaxValue = maxValue - port
+  if distanceToMaxValue < offset:
+    return minValue + offset - distanceToMaxValue - 1
+  return port + offset
+
+proc subtractOffset(port: uint16, offset: uint16, minValue = 1024'u16,
+                    maxValue = uint16.high): uint16 =
+  assert(port >= minValue)
+  assert(port <= maxValue)
+  let distanceToMinValue = port - minValue
+  if distanceToMinValue < offset:
+    return maxValue - offset + distanceToMinValue + 1
+  return port - offset
+
+proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
+  if probedPorts.len == 0:
+    # No probed ports, so our only guess can be that the NAT is a cone-type NAT
+    # and the port mapping preserves the local Port.
+    return @[localPort]
+  let localPortUint = localPort.uint16
+  let probedPortsUint = probedPorts.map(toUint16)
+  if probedPorts.len == 1:
+    # Only one server was used for probing, so we cannot know if the NAT is
+    # symmetric or not. We are trying the probed port (assuming cone-type NAT)
+    # and the next port in a progressive sequence if applicable (assuming
+    # symmetric NAT with progressive port mapping).
+    result.add(probedPorts[0])
+    if probedPortsUint[0] > localPortUint:
+      let offset = probedPortsUint[0] - localPortUint
+      result.add(Port(probedPortsUint[0].addOffset(offset)))
+    elif probedPortsUint[0] < localPortUint:
+      let offset = localPortUint - probedPortsUint[0]
+      result.add(Port(probedPortsUint[0].subtractOffset(offset)))
+    return
+  let deduplicatedPorts = probedPortsUint.deduplicate()
+  if deduplicatedPorts.len() == 1:
+    # It looks like the NAT is a cone-type NAT.
+    return deduplicatedPorts.map(toPort)
+  let probedPortsSorted = probedPortsUint.sorted()
+  let minPort = probedPortsSorted[probedPortsSorted.minIndex()]
+  let maxPort = probedPortsSorted[probedPortsSorted.maxIndex()]
+  var minDistance = uint16.high()
+  var maxDistance = uint16.low()
+  for i in 1 .. probedPortsSorted.len() - 1:
+    # FIXME: use rotated distance
+    let distance = probedPortsSorted[i] - probedPortsSorted[i - 1]
+    minDistance = min(minDistance, distance)
+    maxDistance = max(maxDistance, distance)
+  if maxDistance < 10:
+    if probedPortsUint.isSorted(Ascending):
+      # assume symmetric NAT with positive-progressive port mapping
+      if minDistance == maxDistance:
+        return @[Port(maxPort.addOffset(maxDistance))]
+      else:
+        for i in countup(0'u16, maxDistance):
+          result.add(Port(minPort.addOffset(i)))
+        return
+    if probedPortsUint.isSorted(Descending):
+      # assume symmetric NAT with negative-progressive port mapping
+      if minDistance == maxDistance:
+        return @[Port(minPort.subtractOffset(maxDistance))]
+      else:
+        for i in countup(0'u16, maxDistance):
+          result.add(Port(maxPort.subtractOffset(i)))
+        return
+  # assume symmetric NAT with random port mapping
+  let portRange = maxPort - minPort
+  let first = if portRange > RandomPortCount:
+    minPort
+  else:
+    let notCovered = RandomPortCount - portRange
+    max(minPort - notCovered shr 1, 1024)
+  let last = first + RandomPortCount
+  for i in first .. last:
+    result.add(Port(i))
+
+suite "port prediction tests":
+  test "single port":
+    let predicted = predictPortRange(Port(1234), @[])
+    check(predicted == @[Port(1234)])
+
+  test "single probe equal":
+    let predicted = predictPortRange(Port(1234), @[Port(1234)])
+    check(predicted == @[Port(1234)])
+
+  test "single probe positive-progressive":
+    let predicted = predictPortRange(Port(1234), @[Port(1236)])
+    check(predicted == @[Port(1236), Port(1238)])
+
+  test "single probe negative-progressive":
+    let predicted = predictPortRange(Port(1234), @[Port(1232)])
+    check(predicted == @[Port(1232), Port(1230)])
+
+  test "all equal":
+    let predicted = predictPortRange(Port(1234), @[Port(1234), Port(1234)])
+    check(predicted == @[Port(1234)])
+
+  test "positive-progressive, offset 1":
+    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2035)])
+    check(predicted == @[Port(2036)])
+
+  test "positive-progressive, offset 9":
+    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2043)])
+    check(predicted == @[Port(2052)])
+
+  test "negative-progressive, offset 1":
+    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1099)])
+    check(predicted == @[Port(1098)])
+
+  test "negative-progressive, offset 9":
+    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1091)])
+    check(predicted == @[Port(1082)])
+
+  test "positive-progressive, 3 probed ports, low offset":
+    let predicted = predictPortRange(Port(1234), @[Port(2000), Port(2000), Port(2002)])
+    check(predicted == @[Port(2000), Port(2001), Port(2002)])
+
+  test "negative-progressive, 3 probed ports, low offset":
+    let predicted = predictPortRange(Port(1234), @[Port(2002), Port(2000), Port(2000)])
+    check(predicted == @[Port(2002), Port(2001), Port(2000)])
+
+  test "high port, positive-progressive, offset 1":
+    let predicted = predictPortRange(Port(1234), @[Port(65534), Port(65535)])
+    check(predicted == @[Port(1024)])
+
+  test "high port, positive-progressive, offset 9":
+    let predicted = predictPortRange(Port(1234), @[Port(65520), Port(65529)])
+    check(predicted == @[Port(1026)])
+
+  test "low port, negative-progressive, offset 1":
+    let predicted = predictPortRange(Port(1234), @[Port(1025), Port(1024)])
+    check(predicted == @[Port(65535)])
+
+  test "low port, negative-progressive, offset 9":
+    let predicted = predictPortRange(Port(1234), @[Port(1039), Port(1030)])
+    check(predicted == @[Port(65533)])
+
+  test "random mapping, distance > RandomPortCount":
+    let predicted = predictPortRange(Port(1234), @[Port(3546), Port(7624)])
+    check(predicted == toSeq(countup(3546'u16, 3546'u16 + RandomPortCount)).map(toPort))
+
+  test "random mapping, distance < RandomPortCount":
+    let centerPort = 30000'u16
+    let minPort = centerPort - RandomPortCount.uint16 shr 1 + 1
+    let maxPort = centerPort + RandomPortCount.uint16 shr 1 - 1
+    let predicted = predictPortRange(Port(centerPort), @[Port(minPort), Port(maxPort)])
+    check(predicted == toSeq(countup(minPort - 1, maxPort + 1)).map(toPort))