From f80278aeeed1163e4249b66c8ff0d953857952f4 Mon Sep 17 00:00:00 2001
From: andri lim <jangko128@yahoo.com>
Date: Tue, 15 Jun 2021 02:41:31 +0700
Subject: [PATCH] implement frame reader for buffered payload (#55)

now the frame can choose to read from buffered
payload if available or read from asyncstream
if the buffer is empty
---
 ws/frame.nim   | 29 +++++++++++++++++++++++++++++
 ws/session.nim | 12 +-----------
 ws/types.nim   |  1 +
 3 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/ws/frame.nim b/ws/frame.nim
index a0f52a0a..afe9fde9 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 289d4e1b..a25b9ffc 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 7e95f6c3..01ab81ab 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]