From edfbb169c84f80caa0a1985ef0b5e19f006f8532 Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Wed, 18 May 2022 14:37:53 +0200
Subject: [PATCH] Fix some issues with recv hanging (#109)

* Crash on concurrent reads
* Fix some issues with recv hanging
---
 websock/frame.nim   |  3 +++
 websock/session.nim | 11 ++++++++---
 websock/types.nim   |  1 +
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/websock/frame.nim b/websock/frame.nim
index 707516b..e4d0342 100644
--- a/websock/frame.nim
+++ b/websock/frame.nim
@@ -205,6 +205,9 @@ proc decode*(
 
   frame.length = finalLen
 
+  if frame.length > WSMaxMessageSize:
+    raise newException(WSPayloadLengthError, "Frame too big: " & $frame.length)
+
   # Do we need to apply mask?
   frame.mask = (b1 and 0x80) == 0x80
   if masked == frame.mask:
diff --git a/websock/session.nim b/websock/session.nim
index 6349964..5456f37 100644
--- a/websock/session.nim
+++ b/websock/session.nim
@@ -312,6 +312,10 @@ proc recv*(
   ## Use this to stream data from frames
   ##
 
+  doAssert ws.reading == false, "Only one concurrent read allowed"
+  ws.reading = true
+  defer: ws.reading = false
+
   var consumed = 0
   var pbuffer = cast[ptr UncheckedArray[byte]](data)
   try:
@@ -321,7 +325,8 @@ proc recv*(
 
     while consumed < size:
       if isNil(ws.frame):
-        trace "Empty frame, breaking"
+        assert ws.readyState == ReadyState.Closed
+        trace "Closed connection, breaking"
         break
 
       logScope:
@@ -347,8 +352,8 @@ proc recv*(
         trace "Reading bytes from frame stream", len
         let read = await ws.frame.read(ws.stream.reader, addr pbuffer[consumed], len)
         if read <= 0:
-          trace "Didn't read any bytes, breaking"
-          break
+          trace "Didn't read any bytes, stopping"
+          raise newException(WSClosedError, "WebSocket is closed!")
 
         trace "Read data from frame", read
         consumed += read
diff --git a/websock/types.nim b/websock/types.nim
index f80102d..e843582 100644
--- a/websock/types.nim
+++ b/websock/types.nim
@@ -96,6 +96,7 @@ type
     stream*: AsyncStream
     frame*: Frame
     first*: bool
+    reading*: bool
     proto*: string
 
   Ext* = ref object of RootObj