change port prediction API: allow puncher to retrieve NAT properties

port_prediction.nim

@@ -3,6 +3,28 @@ import net
 import sequtils
 import unittest
 
+type
+  NatType* = enum
+    Unknown,
+    Cone,
+    SymmetricProgressive,
+    SymmetricRandom
+
+  NatProperties* = object
+    case natType*: NatType
+    of Unknown:
+      guess*: seq[uint16]
+    of Cone:
+      prediction*: uint16
+    of SymmetricProgressive:
+      order*: SortOrder
+      previousPort*: uint16
+      minDistance*: uint16
+      maxDistance*: uint16
+    of SymmetricRandom:
+      minPort*: uint16
+      maxPort*: uint16
+
 const RandomPortCount = 10000'u16
 
 proc min(a, b: uint16): uint16 =
@@ -30,28 +52,29 @@ proc subtractOffset(port: uint16, offset: uint16, minValue = 1024'u16,
     return maxValue - offset + distanceToMinValue + 1
   return port - offset
 
-proc predictPortRange*(localPort: uint16, probedPorts: seq[uint16]): seq[uint16] =
+proc getNatProperties*(localPort: uint16, probedPorts: seq[uint16]):
+                      NatProperties =
   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]
+    return NatProperties(natType: Unknown, guess: @[localPort])
   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])
+    result = NatProperties(natType: Unknown, guess: @[probedPorts[0]])
     if probedPorts[0] > localPort:
       let offset = probedPorts[0] - localPort
-      result.add(probedPorts[0].addOffset(offset))
+      result.guess.add(probedPorts[0].addOffset(offset))
     elif probedPorts[0] < localPort:
       let offset = localPort - probedPorts[0]
-      result.add(probedPorts[0].subtractOffset(offset))
+      result.guess.add(probedPorts[0].subtractOffset(offset))
     return
   let deduplicatedPorts = probedPorts.deduplicate()
   if deduplicatedPorts.len() == 1:
     # It looks like the NAT is a cone-type NAT.
-    return deduplicatedPorts
+    return NatProperties(natType: Cone, prediction: deduplicatedPorts[0])
   let probedPortsSorted = probedPorts.sorted()
   let minPort = probedPortsSorted[probedPortsSorted.minIndex()]
   let maxPort = probedPortsSorted[probedPortsSorted.maxIndex()]
@@ -65,107 +88,213 @@ proc predictPortRange*(localPort: uint16, probedPorts: seq[uint16]): seq[uint16]
   if maxDistance < 10:
     if probedPorts.isSorted(Ascending):
       # assume symmetric NAT with positive-progressive port mapping
-      if minDistance == maxDistance:
+      return NatProperties(natType: SymmetricProgressive,
-        return @[maxPort.addOffset(maxDistance)]
+                           order: Ascending, previousPort: maxPort,
-      else:
+                           minDistance: minDistance, maxDistance: maxDistance)
-        for i in countup(0'u16, maxDistance):
-          result.add(minPort.addOffset(i))
-        return
     if probedPorts.isSorted(Descending):
       # assume symmetric NAT with negative-progressive port mapping
-      if minDistance == maxDistance:
+      return NatProperties(natType: SymmetricProgressive,
-        return @[minPort.subtractOffset(maxDistance)]
+                           order: Descending, previousPort: minPort,
-      else:
+                           minDistance: minDistance, maxDistance: maxDistance)
-        for i in countup(0'u16, maxDistance):
-          result.add(maxPort.subtractOffset(i))
-        return
   # assume symmetric NAT with random port mapping
+  return NatProperties(natType: SymmetricRandom, minPort: minPort,
+                       maxPort: maxPort)
+
+proc getNatProperties*(localPort: Port, probedPorts: seq[Port]): NatProperties =
+  getNatProperties(localPort.uint16, probedPorts.map(toUInt16))
+
+proc predictPortRange*(props: NatProperties): seq[Port] =
+  case props.natType
+  of Unknown:
+    result = props.guess.map(toPort)
+  of Cone:
+    result = @[Port(props.prediction)]
+  of SymmetricProgressive:
+    if props.order == Ascending:
+      if props.minDistance == props.maxDistance:
+        return @[Port(props.previousPort.addOffset(props.maxDistance))]
+      else:
+        let minPort = props.previousPort - props.maxDistance
+        for i in countup(0'u16, props.maxDistance):
+          result.add(Port(minPort.addOffset(i)))
+    else:
+      if props.minDistance == props.maxDistance:
+        return @[Port(props.previousPort.subtractOffset(props.maxDistance))]
+      else:
+        let maxPort = props.previousPort + props.maxDistance
+        for i in countup(0'u16, props.maxDistance):
+          result.add(Port(maxPort.subtractOffset(i)))
+  of SymmetricRandom:
     assert(RandomPortCount mod 2 == 0)
-  let center = minPort + (maxPort - minPort) div 2
+    let center = props.minPort + (props.maxPort - props.minPort) div 2
     let half = RandomPortCount div 2
     let first = if (1024'u16 + half) < center:
       min(center - half, uint16.high - RandomPortCount + 1)
     else:
       1024'u16
-  result = newSeq[uint16](RandomPortCount)
+    result = newSeq[Port](RandomPortCount)
     for i in 0'u16 .. RandomPortCount - 1'u16:
-    result[i] = first + i
+      result[i] = Port(first + i)
-
-proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
-  predictPortRange(localPort.uint16, probedPorts.map(toUInt16)).map(toPort)
 
 suite "port prediction tests":
   test "single port":
-    let predicted = predictPortRange(Port(1234), @[])
+    let props = getNatProperties(1234'u16, @[])
+    check(props.natType == Unknown)
+    check(props.guess == @[1234'u16])
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1234)])
 
   test "single probe equal":
-    let predicted = predictPortRange(Port(1234), @[Port(1234)])
+    let props = getNatProperties(1234'u16, @[1234'u16])
+    check(props.natType == Unknown)
+    check(props.guess == @[1234'u16])
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1234)])
 
   test "single probe positive-progressive":
-    let predicted = predictPortRange(Port(1234), @[Port(1236)])
+    let props = getNatProperties(1234'u16, @[1236'u16])
+    check(props.natType == Unknown)
+    check(props.guess == @[1236'u16, 1238'u16])
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1236), Port(1238)])
 
   test "single probe negative-progressive":
-    let predicted = predictPortRange(Port(1234), @[Port(1232)])
+    let props = getNatProperties(1234'u16, @[1232'u16])
+    check(props.natType == Unknown)
+    check(props.guess == @[1232'u16, 1230'u16])
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1232), Port(1230)])
 
   test "all equal":
-    let predicted = predictPortRange(Port(1234), @[Port(1234), Port(1234)])
+    let props = getNatProperties(1234'u16, @[1234'u16, 1234'u16])
+    check(props.natType == Cone)
+    check(props.prediction == 1234'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1234)])
 
   test "positive-progressive, offset 1":
-    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2035)])
+    let props = getNatProperties(1234'u16, @[2034'u16, 2035'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Ascending)
+    check(props.previousPort == 2035'u16)
+    check(props.minDistance == 1'u16)
+    check(props.maxDistance == 1'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(2036)])
 
   test "positive-progressive, offset 9":
-    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2043)])
+    let props = getNatProperties(1234'u16, @[2034'u16, 2043'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Ascending)
+    check(props.previousPort == 2043'u16)
+    check(props.minDistance == 9'u16)
+    check(props.maxDistance == 9'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(2052)])
 
   test "negative-progressive, offset 1":
-    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1099)])
+    let props = getNatProperties(1234'u16, @[1100'u16, 1099'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Descending)
+    check(props.previousPort == 1099'u16)
+    check(props.minDistance == 1'u16)
+    check(props.maxDistance == 1'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1098)])
 
   test "negative-progressive, offset 9":
-    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1091)])
+    let props = getNatProperties(1234'u16, @[1100'u16, 1091'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Descending)
+    check(props.previousPort == 1091'u16)
+    check(props.minDistance == 9'u16)
+    check(props.maxDistance == 9'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1082)])
 
   test "positive-progressive, 3 probed ports, low offset":
-    let predicted = predictPortRange(Port(1234), @[Port(2000), Port(2000), Port(2002)])
+    let props = getNatProperties(1234'u16, @[2000'u16, 2000'u16, 2002'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Ascending)
+    check(props.previousPort == 2002)
+    check(props.minDistance == 0'u16)
+    check(props.maxDistance == 2'u16)
+    let predicted = predictPortRange(props)
     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)])
+    let props = getNatProperties(1234'u16, @[2002'u16, 2000'u16, 2000'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Descending)
+    check(props.previousPort == 2000)
+    check(props.minDistance == 0'u16)
+    check(props.maxDistance == 2'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(2002), Port(2001), Port(2000)])
 
   test "high port, positive-progressive, offset 1":
-    let predicted = predictPortRange(Port(1234), @[Port(65534), Port(65535)])
+    let props = getNatProperties(1234'u16, @[65534'u16, 65535'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Ascending)
+    check(props.previousPort == 65535'u16)
+    check(props.minDistance == 1'u16)
+    check(props.maxDistance == 1'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1024)])
 
   test "high port, positive-progressive, offset 9":
-    let predicted = predictPortRange(Port(1234), @[Port(65520), Port(65529)])
+    let props = getNatProperties(1234'u16, @[65520'u16, 65529'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Ascending)
+    check(props.previousPort == 65529'u16)
+    check(props.minDistance == 9'u16)
+    check(props.maxDistance == 9'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(1026)])
 
   test "low port, negative-progressive, offset 1":
-    let predicted = predictPortRange(Port(1234), @[Port(1025), Port(1024)])
+    let props = getNatProperties(1234'u16, @[1025'u16, 1024'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Descending)
+    check(props.previousPort == 1024'u16)
+    check(props.minDistance == 1'u16)
+    check(props.maxDistance == 1'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(65535)])
 
   test "low port, negative-progressive, offset 9":
-    let predicted = predictPortRange(Port(1234), @[Port(1039), Port(1030)])
+    let props = getNatProperties(1234'u16, @[1039'u16, 1030'u16])
+    check(props.natType == SymmetricProgressive)
+    check(props.order == Descending)
+    check(props.previousPort == 1030'u16)
+    check(props.minDistance == 9'u16)
+    check(props.maxDistance == 9'u16)
+    let predicted = predictPortRange(props)
     check(predicted == @[Port(65533)])
 
   test "random mapping":
-    let predicted = predictPortRange(Port(1234), @[Port(20000), Port(24000)])
+    let props = getNatProperties(1234'u16, @[20000'u16, 24000'u16])
+    check(props.natType == SymmetricRandom)
+    check(props.minPort == 20000'u16)
+    check(props.maxPort == 24000'u16)
+    let predicted = predictPortRange(props)
     let half = RandomPortCount div 2'u16
     check(predicted == toSeq(countup(22000'u16 - half, 22000'u16 + half - 1)).map(toPort))
 
   test "random mapping, low":
-    let predicted = predictPortRange(Port(1234), @[Port(1200), Port(1600)])
+    let props = getNatProperties(1234'u16, @[1200'u16, 1600'u16])
+    check(props.natType == SymmetricRandom)
+    check(props.minPort == 1200'u16)
+    check(props.maxPort == 1600'u16)
+    let predicted = predictPortRange(props)
     check(predicted.len == RandomPortCount.int)
     check(predicted == toSeq(countup(1024'u16, 1024'u16 + RandomPortCount - 1)).map(toPort))
 
   test "random mapping, high":
-    let predicted = predictPortRange(Port(1234), @[Port(65000), Port(65400)])
+    let props = getNatProperties(1234'u16, @[65000'u16, 65400'u16])
+    check(props.natType == SymmetricRandom)
+    check(props.minPort == 65000'u16)
+    check(props.maxPort == 65400'u16)
+    let predicted = predictPortRange(props)
     check(predicted.len == RandomPortCount.int)
     check(predicted == toSeq(countup(uint16.high - RandomPortCount + 1, uint16.high)).map(toPort))

puncher.nim

@@ -42,7 +42,8 @@ proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
            peerProbedPorts: seq[Port], lowTTL: bool, msg: string):
           Future[Attempt] {.async.} =
   let punchFuture = newFuture[Port]("punch")
-  let predictedDstPorts = predictPortRange(peerPort, peerProbedPorts)
+  let natProps = getNatProperties(peerPort, peerProbedPorts)
+  let predictedDstPorts = predictPortRange(natProps)
   let (_, myPort) = puncher.sock.getLocalAddr()
   result = Attempt(srcPort: myPort, dstIp: peerIp, dstPorts: predictedDstPorts,
                    future: punchFuture)