diff --git a/port_prediction.nim b/port_prediction.nim index c291ca9..6b63659 100644 --- a/port_prediction.nim +++ b/port_prediction.nim @@ -1,10 +1,9 @@ import algorithm import net -import random import sequtils import unittest -const RandomPortCount = 10000 +const RandomPortCount = 10000'u16 proc min(a, b: uint16): uint16 = min(a.int32, b.int32).uint16 @@ -81,11 +80,13 @@ proc predictPortRange*(localPort: uint16, probedPorts: seq[uint16]): seq[uint16] result.add(maxPort.subtractOffset(i)) return # assume symmetric NAT with random port mapping - randomize() - let first = if maxPort <= uint16.high - RandomPortCount - 1000'u16: - maxPort + 1000'u16 + assert(RandomPortCount mod 2 == 0) + let center = minPort + (maxPort - minPort) div 2 + let half = RandomPortCount div 2 + let first = if (1024'u16 + half) < center: + min(center - half, uint16.high - RandomPortCount + 1) else: - uint16.high - RandomPortCount + 1024'u16 result = newSeq[uint16](RandomPortCount) for i in 0'u16 .. RandomPortCount - 1'u16: result[i] = first + i @@ -155,5 +156,16 @@ suite "port prediction tests": check(predicted == @[Port(65533)]) test "random mapping": - let predicted = predictPortRange(Port(1234), @[Port(3546), Port(7624)]) - check(predicted.len == RandomPortCount) + let predicted = predictPortRange(Port(1234), @[Port(20000), Port(24000)]) + 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)]) + 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)]) + check(predicted.len == RandomPortCount.int) + check(predicted == toSeq(countup(uint16.high - RandomPortCount + 1, uint16.high)).map(toPort))