From 3696e3f3a5b938e478e473a6089bf8de386d2f04 Mon Sep 17 00:00:00 2001
From: diegomrsantos <diego@status.im>
Date: Tue, 6 Jun 2023 22:57:40 +0200
Subject: [PATCH] Handle cancellation in close (#143)

* Handle cancellation in close

* Using a future instead of sleepAsync
---
 tests/testwebsockets.nim | 31 +++++++++++++++++++++++++++++++
 websock/session.nim      |  5 +++++
 2 files changed, 36 insertions(+)

diff --git a/tests/testwebsockets.nim b/tests/testwebsockets.nim
index a65cb76499..70188e5ee2 100644
--- a/tests/testwebsockets.nim
+++ b/tests/testwebsockets.nim
@@ -170,6 +170,37 @@ suite "Test transmission":
     check string.fromBytes(clientRes) == testString
     await waitForClose(session)
 
+  asyncTest "Close handle cancellation":
+    let testString = "Hello!"
+    let cancelSignal = newFuture[void]()
+
+    proc handle(request: HttpRequest) {.async.} =
+      check request.uri.path == WSPath
+
+      let server = WSServer.new(protos = ["proto"])
+      let ws = await server.handleRequest(request)
+      let servRes = await ws.recvMsg()
+
+      check string.fromBytes(servRes) == testString
+      await ws.waitForClose()
+
+    server = createServer(
+      address = address,
+      handler = handle,
+      flags = {ReuseAddr})
+
+    proc client() {.async, gcsafe.} =
+      let session = await connectClient()
+      await session.send(testString)
+      cancelSignal.complete()
+      expect CancelledError:
+        await session.close()
+
+    let task = client()
+
+    await cancelSignal
+    await task.cancelAndWait()
+
 suite "Test ping-pong":
   setup:
     var
diff --git a/websock/session.nim b/websock/session.nim
index bda209a859..5fc8265442 100644
--- a/websock/session.nim
+++ b/websock/session.nim
@@ -518,11 +518,16 @@ proc close*(
     try:
       while ws.readyState != ReadyState.Closed:
         discard await ws.readFrame()
+    except CancelledError as exc:
+      raise exc
     except CatchableError as exc:
       discard # most likely EOF
   try:
     ws.readyState = ReadyState.Closing
     await gentleCloser(ws, prepareCloseBody(code, reason)).wait(10.seconds)
+  except CancelledError as exc:
+    trace "Cancellation when closing!", exc = exc.msg
+    raise exc
   except CatchableError as exc:
     trace "Exception closing", exc = exc.msg
   finally: