diff --git a/websock/frame.nim b/websock/frame.nim index 707516bc..e4d03420 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 6349964f..5456f379 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 f80102d2..e8435823 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