From eaa72dcdbe26d8236c64711e32cf0c2f8ac65e64 Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Thu, 17 Mar 2022 10:16:48 +0100
Subject: [PATCH] WS Accept timeout (#699)

* Add timeout to WS accept
* Handle more WS errors
---
 .pinned                           |  4 +--
 libp2p/transports/wstransport.nim | 49 +++++++++++++++++++------------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/.pinned b/.pinned
index de0146561..252ceced8 100644
--- a/.pinned
+++ b/.pinned
@@ -13,5 +13,5 @@ serialization;https://github.com/status-im/nim-serialization@#37bc0db558d8571196
 stew;https://github.com/status-im/nim-stew@#bb705bf17b46d2c8f9bfb106d9cc7437009a2501
 testutils;https://github.com/status-im/nim-testutils@#aa6e5216f4b4ab5aa971cdcdd70e1ec1203cedf2
 unittest2;https://github.com/status-im/nim-unittest2@#4e2893eacb916c7678fdc4935ff7420f13bf3a9c
-websock;https://github.com/status-im/nim-websock@#73edde4417f7b45003113b7a34212c3ccd95b9fd
-zlib;https://github.com/status-im/nim-zlib@#74cdeb54b21bededb5a515d36f608bc1850555a2
\ No newline at end of file
+websock;https://github.com/status-im/nim-websock@#853299e399746eff4096870067cbc61861ecd534
+zlib;https://github.com/status-im/nim-zlib@#74cdeb54b21bededb5a515d36f608bc1850555a2
diff --git a/libp2p/transports/wstransport.nim b/libp2p/transports/wstransport.nim
index 4f6444294..2dbdd1c44 100644
--- a/libp2p/transports/wstransport.nim
+++ b/libp2p/transports/wstransport.nim
@@ -30,6 +30,8 @@ export transport, websock
 const
   WsTransportTrackerName* = "libp2p.wstransport"
 
+  DefaultHeadersTimeout = 3.seconds
+
 type
   WsStream = ref object of Connection
     session: WSSession
@@ -92,6 +94,7 @@ type
     tlsCertificate: TLSCertificate
     tlsFlags: set[TLSFlags]
     flags: set[ServerFlags]
+    handshakeTimeout: Duration
     factories: seq[ExtFactory]
     rng: Rng
 
@@ -131,9 +134,13 @@ method start*(
           address = ma.initTAddress().tryGet(),
           tlsPrivateKey = self.tlsPrivateKey,
           tlsCertificate = self.tlsCertificate,
-          flags = self.flags)
+          flags = self.flags,
+          handshakeTimeout = self.handshakeTimeout)
       else:
-        HttpServer.create(ma.initTAddress().tryGet())
+        HttpServer.create(
+          ma.initTAddress().tryGet(),
+          handshakeTimeout = self.handshakeTimeout
+        )
 
     self.httpservers &= httpserver
 
@@ -222,19 +229,19 @@ method accept*(self: WsTransport): Future[Connection] {.async, gcsafe.} =
   if not self.running:
     raise newTransportClosedError()
 
+  if self.acceptFuts.len <= 0:
+    self.acceptFuts = self.httpservers.mapIt(it.accept())
+
+  if self.acceptFuts.len <= 0:
+    return
+
+  let
+    finished = await one(self.acceptFuts)
+    index = self.acceptFuts.find(finished)
+
+  self.acceptFuts[index] = self.httpservers[index].accept()
+
   try:
-    if self.acceptFuts.len <= 0:
-      self.acceptFuts = self.httpservers.mapIt(it.accept())
-
-    if self.acceptFuts.len <= 0:
-      return
-
-    let
-      finished = await one(self.acceptFuts)
-      index = self.acceptFuts.find(finished)
-
-    self.acceptFuts[index] = self.httpservers[index].accept()
-
     let req = await finished
 
     try:
@@ -250,6 +257,8 @@ method accept*(self: WsTransport): Future[Connection] {.async, gcsafe.} =
     debug "OS Error", exc = exc.msg
   except WebSocketError as exc:
     debug "Websocket Error", exc = exc.msg
+  except HttpError as exc:
+    debug "Http Error", exc = exc.msg
   except AsyncStreamError as exc:
     debug "AsyncStream Error", exc = exc.msg
   except TransportTooManyError as exc:
@@ -301,7 +310,8 @@ proc new*(
   tlsFlags: set[TLSFlags] = {},
   flags: set[ServerFlags] = {},
   factories: openArray[ExtFactory] = [],
-  rng: Rng = nil): T =
+  rng: Rng = nil,
+  handshakeTimeout = DefaultHeadersTimeout): T =
 
   T(
     upgrader: upgrade,
@@ -310,14 +320,16 @@ proc new*(
     tlsFlags: tlsFlags,
     flags: flags,
     factories: @factories,
-    rng: rng)
+    rng: rng,
+    handshakeTimeout: handshakeTimeout)
 
 proc new*(
   T: typedesc[WsTransport],
   upgrade: Upgrade,
   flags: set[ServerFlags] = {},
   factories: openArray[ExtFactory] = [],
-  rng: Rng = nil): T =
+  rng: Rng = nil,
+  handshakeTimeout = DefaultHeadersTimeout): T =
 
   T.new(
     upgrade = upgrade,
@@ -325,4 +337,5 @@ proc new*(
     tlsCertificate = nil,
     flags = flags,
     factories = @factories,
-    rng = rng)
+    rng = rng,
+    handshakeTimeout = handshakeTimeout)