diff --git a/ws/frame.nim b/ws/frame.nim index a0f52a0..afe9fde 100644 --- a/ws/frame.nim +++ b/ws/frame.nim @@ -56,6 +56,35 @@ proc mask*( template remainder*(frame: Frame): uint64 = frame.length - frame.consumed +proc read*(frame: Frame, + reader: AsyncStreamReader, + pbytes: pointer, + nbytes: int): Future[int] {.async.} = + + # read data from buffered payload if available + # e.g. data processed by extensions + var readLen = 0 + if frame.offset < frame.data.len: + readLen = min(frame.data.len - frame.offset, nbytes) + copyMem(pbytes, addr frame.data[frame.offset], readLen) + frame.offset += readLen + + var pbuf = cast[ptr UncheckedArray[byte]](pbytes) + if readLen < nbytes: + let len = min(nbytes - readLen, frame.remainder.int - readLen) + readLen += await reader.readOnce(addr pbuf[readLen], len) + + if frame.mask and readLen > 0: + # unmask data using offset + mask( + pbuf.toOpenArray(0, readLen - 1), + frame.maskKey, + frame.consumed.int + ) + + frame.consumed += readLen.uint64 + return readLen + proc encode*( frame: Frame, extensions: seq[Ext] = @[]): Future[seq[byte]] {.async.} = diff --git a/ws/session.nim b/ws/session.nim index 289d4e1..a25b9ff 100644 --- a/ws/session.nim +++ b/ws/session.nim @@ -359,23 +359,13 @@ proc recv*( let len = min(ws.frame.remainder.int, size - consumed) if len > 0: trace "Reading bytes from frame stream", len - let read = await ws.stream.reader.readOnce(addr pbuffer[consumed], 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 "Read data from frame", read - - if ws.frame.mask: - trace "Unmasking frame" - # unmask data using offset - mask( - pbuffer.toOpenArray(consumed, (consumed + read) - 1), - ws.frame.maskKey, - ws.frame.consumed.int) - consumed += read - ws.frame.consumed += read.uint64 # all has been consumed from the frame # read the next frame diff --git a/ws/types.nim b/ws/types.nim index 7e95f6c..01ab81a 100644 --- a/ws/types.nim +++ b/ws/types.nim @@ -61,6 +61,7 @@ type maskKey*: MaskKey ## Masking key length*: uint64 ## Message size. consumed*: uint64 ## how much has been consumed from the frame + offset*: int ## offset of buffered payload data StatusCodes* = distinct range[0..4999]