From c43aacdc81ea044518b24e82ff942671314e3dc8 Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Tue, 22 Nov 2022 10:44:08 +0100
Subject: [PATCH] WS: add timeout to handleRequest (#796)

---
 libp2p/transports/wstransport.nim | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/libp2p/transports/wstransport.nim b/libp2p/transports/wstransport.nim
index 9ded1c1e7..220a21b45 100644
--- a/libp2p/transports/wstransport.nim
+++ b/libp2p/transports/wstransport.nim
@@ -35,14 +35,18 @@ logScope:
 export transport, websock, results
 
 const
-  WsTransportTrackerName* = "libp2p.wstransport"
-
   DefaultHeadersTimeout = 3.seconds
 
 type
   WsStream = ref object of Connection
     session: WSSession
 
+method initStream*(s: WsStream) =
+  if s.objName.len == 0:
+    s.objName = "WsStream"
+
+  procCall Connection(s).initStream()
+
 proc new*(T: type WsStream,
            session: WSSession,
            dir: Direction,
@@ -129,7 +133,7 @@ method start*(
     factories = self.factories,
     rng = self.rng)
 
-  
+
   for i, ma in addrs:
     let isWss =
       if WSS.match(ma):
@@ -256,7 +260,7 @@ method accept*(self: WsTransport): Future[Connection] {.async, gcsafe.} =
 
     try:
       let
-        wstransp = await self.wsserver.handleRequest(req)
+        wstransp = await self.wsserver.handleRequest(req).wait(self.handshakeTimeout)
         isSecure = self.httpservers[index].secure
 
       return await self.connHandler(wstransp, isSecure, Direction.In)
@@ -273,6 +277,8 @@ method accept*(self: WsTransport): Future[Connection] {.async, gcsafe.} =
     debug "AsyncStream Error", exc = exc.msg
   except TransportTooManyError as exc:
     debug "Too many files opened", exc = exc.msg
+  except AsyncTimeoutError as exc:
+    debug "Timed out", exc = exc.msg
   except TransportUseClosedError as exc:
     debug "Server was closed", exc = exc.msg
     raise newTransportClosedError(exc)