Add port number argument.

Add some documentation comments. Fix tests.
status-im · Mar 11, 2024 · 9101e99 · 9101e99
1 parent 98e7004
commit 9101e99
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 51 deletions.
diff --git a/chronos/transports/common.nim b/chronos/transports/common.nim
@@ -784,8 +784,11 @@ proc setDualstack*(socket: AsyncFD,
       ? getDomain(socket)
   setDualstack(socket, family, flag)
 
-proc getAutoAddress*(): TransportAddress =
-  if isAvailable(AddressFamily.IPv6):
-    AnyAddress6
-  else:
-    AnyAddress
+proc getAutoAddress*(port: Port): TransportAddress =
+  var res =
+    if isAvailable(AddressFamily.IPv6):
+      AnyAddress6
+    else:
+      AnyAddress
+  res.port = port
+  res
diff --git a/chronos/transports/datagram.nim b/chronos/transports/datagram.nim
@@ -824,31 +824,46 @@ proc newDatagramTransport6*[T](cbproc: UnsafeDatagramCallback,
                              cast[pointer](udata), child, bufSize, ttl,
                              dualstack)
 
-proc newAutoDatagramTransport*(cbproc: DatagramCallback,
-                               flags: set[ServerFlags] = {},
-                               udata: pointer = nil,
-                               child: DatagramTransport = nil,
-                               bufSize: int = DefaultDatagramBufferSize,
-                               ttl: int = 0,
-                               dualstack = DualStackType.Auto
-                              ): DatagramTransport {.
+proc newDatagramTransport*(cbproc: DatagramCallback,
+                           port: Port,
+                           flags: set[ServerFlags] = {},
+                           udata: pointer = nil,
+                           child: DatagramTransport = nil,
+                           bufSize: int = DefaultDatagramBufferSize,
+                           ttl: int = 0,
+                           dualstack = DualStackType.Auto
+                          ): DatagramTransport {.
      raises: [TransportOsError].} =
-  let host = getAutoAddress()
+  ## Create new UDP datagram transport (IPv6) and bind it to ANY_ADDRESS.
+  ## Depending on OS settings procedure perform an attempt to create transport
+  ## using IPv6 ANY_ADDRESS, if its not available it will try to bind transport
+  ## to IPv4 ANY_ADDRESS.
+  ##
+  ## ``cbproc`` - callback which will be called, when new datagram received.
+  ## ``port`` - port number.
+  ## ``sock`` - application-driven socket to use.
+  ## ``flags`` - flags that will be applied to socket.
+  ## ``udata`` - custom argument which will be passed to ``cbproc``.
+  ## ``bufSize`` - size of internal buffer.
+  ## ``ttl`` - TTL for UDP datagram packet (only usable when flags has
+  ## ``Broadcast`` option).
+  let host = getAutoAddress(port)
   newDatagramTransportCommon(cbproc, host, host, asyncInvalidSocket, flags,
                              cast[pointer](udata), child, bufSize, ttl,
                              dualstack)
 
-proc newAutoDatagramTransport*[T](cbproc: DatagramCallback,
-                                  flags: set[ServerFlags] = {},
-                                  udata: ref T,
-                                  child: DatagramTransport = nil,
-                                  bufSize: int = DefaultDatagramBufferSize,
-                                  ttl: int = 0,
-                                  dualstack = DualStackType.Auto
-                                 ): DatagramTransport {.
+proc newDatagramTransport*[T](cbproc: DatagramCallback,
+                              port: Port,
+                              flags: set[ServerFlags] = {},
+                              udata: ref T,
+                              child: DatagramTransport = nil,
+                              bufSize: int = DefaultDatagramBufferSize,
+                              ttl: int = 0,
+                              dualstack = DualStackType.Auto
+                             ): DatagramTransport {.
      raises: [TransportOsError].} =
   let
-    host = getAutoAddress()
+    host = getAutoAddress(port)
     fflags = flags + {GCUserData}
   GC_ref(udata)
   newDatagramTransportCommon(cbproc, host, host, asyncInvalidSocket, fflags,

diff --git a/chronos/transports/stream.nim b/chronos/transports/stream.nim
@@ -2130,7 +2130,8 @@ proc createStreamServer*(host: TransportAddress,
   createStreamServer(host, StreamCallback2(nil), flags, sock, backlog, bufferSize,
                      child, init, cast[pointer](udata), dualstack)
 
-proc createStreamServer*(flags: set[ServerFlags] = {},
+proc createStreamServer*(port = Port(0),
+                         flags: set[ServerFlags] = {},
                          sock: AsyncFD = asyncInvalidSocket,
                          backlog: int = DefaultBacklogSize,
                          bufferSize: int = DefaultStreamBufferSize,
@@ -2141,11 +2142,12 @@ proc createStreamServer*(flags: set[ServerFlags] = {},
     raises: [TransportOsError].} =
   ## Create stream server which will be bound to IPv6 address `::`, if IPv6
   ## available, and bound to IPv4 address `0.0.0.0`, if IPv6 is not available.
-  createStreamServer(getAutoAddress(), StreamCallback2(nil), flags, sock,
+  createStreamServer(getAutoAddress(port), StreamCallback2(nil), flags, sock,
                      backlog, bufferSize, child, init, cast[pointer](udata),
                      dualstack)
 
 proc createStreamServer*(cbproc: StreamCallback2,
+                         port = Port(0),
                          flags: set[ServerFlags] = {},
                          sock: AsyncFD = asyncInvalidSocket,
                          backlog: int = DefaultBacklogSize,
@@ -2157,8 +2159,8 @@ proc createStreamServer*(cbproc: StreamCallback2,
     raises: [TransportOsError].} =
   ## Create stream server which will be bound to IPv6 address `::`, if IPv6
   ## available, and bound to IPv4 address `0.0.0.0`, if IPv6 is not available.
-  createStreamServer(getAutoAddress(), cbproc, flags, sock, backlog, bufferSize,
-                     child, init, cast[pointer](udata), dualstack)
+  createStreamServer(getAutoAddress(port), cbproc, flags, sock, backlog,
+                     bufferSize, child, init, cast[pointer](udata), dualstack)
 
 proc createStreamServer*[T](host: TransportAddress,
                             cbproc: StreamCallback2,
@@ -2209,6 +2211,7 @@ proc createStreamServer*[T](host: TransportAddress,
                      child, init, cast[pointer](udata), dualstack)
 
 proc createStreamServer*[T](cbproc: StreamCallback2,
+                            port: Port,
                             flags: set[ServerFlags] = {},
                             udata: ref T,
                             sock: AsyncFD = asyncInvalidSocket,
@@ -2222,10 +2225,11 @@ proc createStreamServer*[T](cbproc: StreamCallback2,
   ## available, and bound to IPv4 address `0.0.0.0`, if IPv6 is not available.
   let fflags = flags + {GCUserData}
   GC_ref(udata)
-  createStreamServer(getAutoAddress(), cbproc, fflags, sock, backlog,
+  createStreamServer(getAutoAddress(port), cbproc, fflags, sock, backlog,
                      bufferSize, child, init, cast[pointer](udata), dualstack)
 
-proc createStreamServer*[T](flags: set[ServerFlags] = {},
+proc createStreamServer*[T](port: Port,
+                            flags: set[ServerFlags] = {},
                             udata: ref T,
                             sock: AsyncFD = asyncInvalidSocket,
                             backlog: int = DefaultBacklogSize,
@@ -2238,7 +2242,7 @@ proc createStreamServer*[T](flags: set[ServerFlags] = {},
   ## available, and bound to IPv4 address `0.0.0.0`, if IPv6 is not available.
   let fflags = flags + {GCUserData}
   GC_ref(udata)
-  createStreamServer(getAutoAddress(), StreamCallback2(nil), fflags, sock,
+  createStreamServer(getAutoAddress(port), StreamCallback2(nil), fflags, sock,
                      backlog, bufferSize, child, init, cast[pointer](udata),
                      dualstack)
 

diff --git a/tests/testdatagram.nim b/tests/testdatagram.nim
@@ -628,7 +628,7 @@ suite "Datagram Transport test suite":
     await allFutures(sdgram.closeWait(), cdgram.closeWait())
     res == 1
 
-  proc performAutoAddressTest(ra: TransportAddress): Future[bool] {.async.} =
+  proc performAutoAddressTest(family: AddressFamily): Future[bool] {.async.} =
     var
       expectStr = "AUTO MESSAGE"
       event = newAsyncEvent()
@@ -653,22 +653,37 @@ suite "Datagram Transport test suite":
          async: (raises: []).} =
       discard
 
+    let sdgram = newDatagramTransport(process1, Port(0))
+
+    var
+      address =
+        case family
+        of AddressFamily.IPv4:
+          initTAddress("127.0.0.1:0")
+        of AddressFamily.IPv6:
+          initTAddress("::1:0")
+        of AddressFamily.Unix, AddressFamily.None:
+          raiseAssert "Not allowed"
+
     let
-      sdgram = newAutoDatagramTransport(process1)
-      cdgram = newAutoDatagramTransport(process2)
+      cdgram =
+        case family
+        of AddressFamily.IPv4:
+          newDatagramTransport(process2, local = address)
+        of AddressFamily.IPv6:
+          newDatagramTransport6(process2, local = address)
+        of AddressFamily.Unix, AddressFamily.None:
+          raiseAssert "Not allowed"
 
-    var address = ra
     address.port = sdgram.localAddress().port
 
-    echo address
-
     try:
       await noCancel cdgram.sendTo(address, addr expectStr[0], len(expectStr))
     except TransportError:
       discard
 
     try:
-      await event.wait().wait(50000.milliseconds)
+      await event.wait().wait(1.seconds)
     except CatchableError:
       discard
 
@@ -780,19 +795,17 @@ suite "Datagram Transport test suite":
   asyncTest "[IP] Auto-address constructor test":
     if isAvailable(AddressFamily.IPv6):
       check:
-        (await performAutoAddressTest(initTAddress("::1:0"))) == true
+        (await performAutoAddressTest(AddressFamily.IPv6)) == true
       # If IPv6 is available newAutoDatagramTransport should bind to `::` - this
       # means that we should be able to connect to it via IPV4_MAPPED address,
       # but only when IPv4 is also available.
       if isAvailable(AddressFamily.IPv4):
         check:
-          (await performAutoAddressTest(
-            initTAddress("127.0.0.1:0").toIPv6())) == true
+          (await performAutoAddressTest(AddressFamily.IPv4)) == true
     else:
       # If IPv6 is not available newAutoDatagramTransport should bind to
       # `0.0.0.0` - this means we should be able to connect to it via IPv4
       # address.
       if isAvailable(AddressFamily.IPv4):
         check:
-          (await performAutoAddressTest(
-            initTAddress("127.0.0.1:0"))) == true
+          (await performAutoAddressTest(AddressFamily.IPv4)) == true
diff --git a/tests/teststream.nim b/tests/teststream.nim
@@ -1486,20 +1486,36 @@ suite "Stream Transport test suite":
     await server.closeWait()
     testResult
 
-  proc performAutoAddressTest(ra: TransportAddress): Future[bool] {.
+  proc performAutoAddressTest(family: AddressFamily): Future[bool] {.
        async: (raises: []).} =
     let server =
       try:
-        createStreamServer()
+        createStreamServer(Port(0))
       except TransportOsError as exc:
         raiseAssert exc.msg
-    var address = ra
+
+    var
+      address =
+        case family
+        of AddressFamily.IPv4:
+          try:
+            initTAddress("127.0.0.1:0")
+          except TransportAddressError as exc:
+            raiseAssert exc.msg
+        of AddressFamily.IPv6:
+          try:
+            initTAddress("::1:0")
+          except TransportAddressError as exc:
+            raiseAssert exc.msg
+        of AddressFamily.Unix, AddressFamily.None:
+          raiseAssert "Not allowed"
+
     address.port = server.localAddress().port
     var acceptFut = server.accept()
     let
       clientTransp =
         try:
-          let res = await connect(address).wait(500.milliseconds)
+          let res = await connect(address).wait(2.seconds)
           Opt.some(res)
         except CatchableError:
           Opt.none(StreamTransport)
@@ -1715,21 +1731,19 @@ suite "Stream Transport test suite":
   asyncTest "[IP] Auto-address constructor test":
     if isAvailable(AddressFamily.IPv6):
       check:
-        (await performAutoAddressTest(initTAddress("::1:0"))) == true
+        (await performAutoAddressTest(AddressFamily.IPv6)) == true
       # If IPv6 is available createStreamServer should bind to `::` this means
       # that we should be able to connect to it via IPV4_MAPPED address, but
       # only when IPv4 is also available.
       if isAvailable(AddressFamily.IPv4):
         check:
-          (await performAutoAddressTest(
-            initTAddress("127.0.0.1:0").toIPv6())) == true
+          (await performAutoAddressTest(AddressFamily.IPv4)) == true
     else:
       # If IPv6 is not available createStreamServer should bind to `0.0.0.0`
       # this means we should be able to connect to it via IPV4 address.
       if isAvailable(AddressFamily.IPv4):
         check:
-          (await performAutoAddressTest(
-            initTAddress("127.0.0.1:0"))) == true
+          (await performAutoAddressTest(AddressFamily.IPv4)) == true
 
   test "File descriptors leak test":
     when defined(windows):