Remove ids from listener when closing connection

quic/connection.nim

@@ -9,6 +9,7 @@ type
     quic: Ngtcp2Connection
     loop: Future[void]
     closed: bool
+    onClose: proc()
   IncomingConnection = ref object of Connection
   OutgoingConnection = ref object of Connection
 
@@ -18,6 +19,9 @@ proc ids*(connection: Connection): seq[ConnectionId] =
 proc `onNewId=`*(connection: Connection, callback: proc (id: ConnectionId)) =
   connection.quic.onNewId = callback
 
+proc `onClose=`*(connection: Connection, callback: proc()) =
+  connection.onClose = callback
+
 proc startSending(connection: Connection, remote: TransportAddress) =
   proc send {.async.} =
     let datagram = await connection.quic.outgoing.get()
@@ -55,12 +59,11 @@ method closeUdp(connection: Connection) {.async, base.} =
 method closeUdp(connection: OutgoingConnection) {.async.} =
   await connection.udp.closeWait()
 
-proc closeQuic(connection: Connection) {.async.} =
-  connection.quic.destroy()
-
 proc close*(connection: Connection) {.async.} =
   if not connection.closed:
     connection.closed = true
     await connection.stopSending()
     await connection.closeUdp()
-    await connection.closeQuic()
+    if connection.onClose != nil:
+      connection.onClose()
+    connection.quic.destroy()

quic/listener.nim

@@ -20,6 +20,9 @@ proc addConnection(listener: Listener, connection: Connection,
                    firstId: ConnectionId) {.async.} =
   connection.onNewId = proc (newId: ConnectionId) =
     listener.connections[newId] = connection
+  connection.onClose = proc =
+    for id in connection.ids & firstId:
+      listener.connections.del(id)
   for id in connection.ids & firstId:
     listener.connections[id] = connection
   await listener.incoming.put(connection)

tests/quic/testListener.nim

@@ -45,3 +45,18 @@ suite "listener":
 
     await first.close()
     await second.close()
+
+  asynctest "forgets connection ids when connection closes":
+    let listener = newListener(address)
+    defer: await listener.stop()
+
+    let datagram = exampleQuicDatagram()
+    await datagram.sendTo(address)
+
+    let first = await listener.waitForIncoming.wait(100.milliseconds)
+    await first.close()
+
+    await datagram.sendTo(address)
+
+    let second = await listener.waitForIncoming.wait(100.milliseconds)
+    await second.close()