From 81a39f32c1e92233e30340be64075bbcf95bf439 Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Wed, 18 Nov 2020 19:32:31 +0100
Subject: [PATCH] move Port -> uint16 conversions to a wrapper proc

---
 port_prediction.nim | 43 ++++++++++++++++++++++---------------------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/port_prediction.nim b/port_prediction.nim
index 5dc8e1e..c291ca9 100644
--- a/port_prediction.nim
+++ b/port_prediction.nim
@@ -9,7 +9,7 @@ const RandomPortCount = 10000
 proc min(a, b: uint16): uint16 =
   min(a.int32, b.int32).uint16
 
-proc toUint16(p: Port): uint16 = uint16(p)
+proc toUInt16(p: Port): uint16 = uint16(p)
 
 proc toPort(u: uint16): Port = Port(u)
 
@@ -31,31 +31,29 @@ proc subtractOffset(port: uint16, offset: uint16, minValue = 1024'u16,
     return maxValue - offset + distanceToMinValue + 1
   return port - offset
 
-proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
+proc predictPortRange*(localPort: uint16, probedPorts: seq[uint16]): seq[uint16] =
   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)))
+    if probedPorts[0] > localPort:
+      let offset = probedPorts[0] - localPort
+      result.add(probedPorts[0].addOffset(offset))
+    elif probedPorts[0] < localPort:
+      let offset = localPort - probedPorts[0]
+      result.add(probedPorts[0].subtractOffset(offset))
     return
-  let deduplicatedPorts = probedPortsUint.deduplicate()
+  let deduplicatedPorts = probedPorts.deduplicate()
   if deduplicatedPorts.len() == 1:
     # It looks like the NAT is a cone-type NAT.
-    return deduplicatedPorts.map(toPort)
-  let probedPortsSorted = probedPortsUint.sorted()
+    return deduplicatedPorts
+  let probedPortsSorted = probedPorts.sorted()
   let minPort = probedPortsSorted[probedPortsSorted.minIndex()]
   let maxPort = probedPortsSorted[probedPortsSorted.maxIndex()]
   var minDistance = uint16.high()
@@ -66,21 +64,21 @@ proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
     minDistance = min(minDistance, distance)
     maxDistance = max(maxDistance, distance)
   if maxDistance < 10:
-    if probedPortsUint.isSorted(Ascending):
+    if probedPorts.isSorted(Ascending):
       # assume symmetric NAT with positive-progressive port mapping
       if minDistance == maxDistance:
-        return @[Port(maxPort.addOffset(maxDistance))]
+        return @[maxPort.addOffset(maxDistance)]
       else:
         for i in countup(0'u16, maxDistance):
-          result.add(Port(minPort.addOffset(i)))
+          result.add(minPort.addOffset(i))
         return
-    if probedPortsUint.isSorted(Descending):
+    if probedPorts.isSorted(Descending):
       # assume symmetric NAT with negative-progressive port mapping
       if minDistance == maxDistance:
-        return @[Port(minPort.subtractOffset(maxDistance))]
+        return @[minPort.subtractOffset(maxDistance)]
       else:
         for i in countup(0'u16, maxDistance):
-          result.add(Port(maxPort.subtractOffset(i)))
+          result.add(maxPort.subtractOffset(i))
         return
   # assume symmetric NAT with random port mapping
   randomize()
@@ -88,9 +86,12 @@ proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
     maxPort + 1000'u16
   else:
     uint16.high - RandomPortCount
-  result = newSeq[Port](RandomPortCount)
+  result = newSeq[uint16](RandomPortCount)
   for i in 0'u16 .. RandomPortCount - 1'u16:
-    result[i] = Port(first + i)
+    result[i] = 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":