diff --git a/port_prediction.nim b/port_prediction.nim index 653102f..35d2f4b 100644 --- a/port_prediction.nim +++ b/port_prediction.nim @@ -25,8 +25,6 @@ type minPort*: uint16 maxPort*: uint16 -const RandomPortCount = 1000'u16 - proc min(a, b: uint16): uint16 = min(a.int32, b.int32).uint16 @@ -103,7 +101,7 @@ proc getNatProperties*(localPort: uint16, probedPorts: seq[uint16]): proc getNatProperties*(localPort: Port, probedPorts: seq[Port]): NatProperties = getNatProperties(localPort.uint16, probedPorts.map(toUInt16)) -proc predictPortRange*(props: NatProperties): seq[Port] = +proc predictPortRange*(props: NatProperties, maxPortCount: uint16): seq[Port] = case props.natType of Unknown: result = props.guess.map(toPort) @@ -125,51 +123,54 @@ proc predictPortRange*(props: NatProperties): seq[Port] = for i in countup(0'u16, props.maxDistance): result.add(Port(maxPort.subtractOffset(i))) of SymmetricRandom: - assert(RandomPortCount mod 2 == 0) + assert(maxPortCount > 0) let center = props.minPort + (props.maxPort - props.minPort) div 2 - let half = RandomPortCount div 2 + let half = maxPortCount div 2 let first = if (1024'u16 + half) < center: - min(center - half, uint16.high - RandomPortCount + 1) + min(center - half, uint16.high - maxPortCount + 1) else: 1024'u16 - result = newSeq[Port](RandomPortCount) - for i in 0'u16 .. RandomPortCount - 1'u16: + result = newSeq[Port](maxPortCount) + for i in 0'u16 .. maxPortCount - 1'u16: result[i] = Port(first + i) suite "port prediction tests": + setup: + let maxPortCount = 1000'u16 + test "single port": let props = getNatProperties(1234'u16, @[]) check(props.natType == Unknown) check(props.guess == @[1234'u16]) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1234)]) test "single probe equal": let props = getNatProperties(1234'u16, @[1234'u16]) check(props.natType == Unknown) check(props.guess == @[1234'u16]) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1234)]) test "single probe positive-progressive": let props = getNatProperties(1234'u16, @[1236'u16]) check(props.natType == Unknown) check(props.guess == @[1236'u16, 1238'u16]) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1236), Port(1238)]) test "single probe negative-progressive": let props = getNatProperties(1234'u16, @[1232'u16]) check(props.natType == Unknown) check(props.guess == @[1232'u16, 1230'u16]) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1232), Port(1230)]) test "all equal": let props = getNatProperties(1234'u16, @[1234'u16, 1234'u16]) check(props.natType == Cone) check(props.prediction == 1234'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1234)]) test "positive-progressive, offset 1": @@ -179,7 +180,7 @@ suite "port prediction tests": check(props.previousPort == 2035'u16) check(props.minDistance == 1'u16) check(props.maxDistance == 1'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(2036)]) test "positive-progressive, offset 9": @@ -189,7 +190,7 @@ suite "port prediction tests": check(props.previousPort == 2043'u16) check(props.minDistance == 9'u16) check(props.maxDistance == 9'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(2052)]) test "negative-progressive, offset 1": @@ -199,7 +200,7 @@ suite "port prediction tests": check(props.previousPort == 1099'u16) check(props.minDistance == 1'u16) check(props.maxDistance == 1'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1098)]) test "negative-progressive, offset 9": @@ -209,7 +210,7 @@ suite "port prediction tests": check(props.previousPort == 1091'u16) check(props.minDistance == 9'u16) check(props.maxDistance == 9'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1082)]) test "positive-progressive, 3 probed ports, low offset": @@ -219,7 +220,7 @@ suite "port prediction tests": check(props.previousPort == 2002) check(props.minDistance == 0'u16) check(props.maxDistance == 2'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(2000), Port(2001), Port(2002)]) test "negative-progressive, 3 probed ports, low offset": @@ -229,7 +230,7 @@ suite "port prediction tests": check(props.previousPort == 2000) check(props.minDistance == 0'u16) check(props.maxDistance == 2'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(2002), Port(2001), Port(2000)]) test "high port, positive-progressive, offset 1": @@ -239,7 +240,7 @@ suite "port prediction tests": check(props.previousPort == 65535'u16) check(props.minDistance == 1'u16) check(props.maxDistance == 1'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1024)]) test "high port, positive-progressive, offset 9": @@ -249,7 +250,7 @@ suite "port prediction tests": check(props.previousPort == 65529'u16) check(props.minDistance == 9'u16) check(props.maxDistance == 9'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(1026)]) test "low port, negative-progressive, offset 1": @@ -259,7 +260,7 @@ suite "port prediction tests": check(props.previousPort == 1024'u16) check(props.minDistance == 1'u16) check(props.maxDistance == 1'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(65535)]) test "low port, negative-progressive, offset 9": @@ -269,7 +270,7 @@ suite "port prediction tests": check(props.previousPort == 1030'u16) check(props.minDistance == 9'u16) check(props.maxDistance == 9'u16) - let predicted = predictPortRange(props) + let predicted = predictPortRange(props, maxPortCount) check(predicted == @[Port(65533)]) test "random port allocation": @@ -277,8 +278,8 @@ suite "port prediction tests": check(props.natType == SymmetricRandom) check(props.minPort == 20000'u16) check(props.maxPort == 24000'u16) - let predicted = predictPortRange(props) - let half = RandomPortCount div 2'u16 + let predicted = predictPortRange(props, maxPortCount) + let half = maxPortCount div 2'u16 check(predicted == toSeq(countup(22000'u16 - half, 22000'u16 + half - 1)).map(toPort)) test "random port allocation": @@ -286,15 +287,15 @@ suite "port prediction tests": 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)) + let predicted = predictPortRange(props, maxPortCount) + check(predicted.len == maxPortCount.int) + check(predicted == toSeq(countup(1024'u16, 1024'u16 + maxPortCount - 1)).map(toPort)) test "random port allocation, high": 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)) + let predicted = predictPortRange(props, maxPortCount) + check(predicted.len == maxPortCount.int) + check(predicted == toSeq(countup(uint16.high - maxPortCount + 1, uint16.high)).map(toPort)) diff --git a/puncher.nim b/puncher.nim index 4afcc0f..f1eb304 100644 --- a/puncher.nim +++ b/puncher.nim @@ -59,7 +59,6 @@ proc initPuncher*(sock: AsyncSocket, probedSrcPorts: seq[Port]): Puncher = proc primarySrcPort*(puncher: Puncher): Port = puncher.socks[0].srcPort -# TODO: lowTTL -> isInitiating, if isInitiating: punch with all auxSocks, else only use 70 proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port, peerProbedPorts: seq[Port], isInitiating: bool, msg: string): Future[Attempt] {.async.} = @@ -78,7 +77,7 @@ proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port, InitiatorMaxSockCount else: ResponderMaxSockCount - let predictedDstPorts = predictPortRange(peerNatProps) + let predictedDstPorts = predictPortRange(peerNatProps, sockCount.uint16) result = Attempt(dstIp: peerIp, dstPorts: predictedDstPorts, future: punchFuture) if puncher.attempts.contains(result):