From 6d3865fd43c2a364520dae849482fae2d386c0cc Mon Sep 17 00:00:00 2001
From: Antonio Nuno Monteiro <anmonteiro@gmail.com>
Date: Sat, 24 Aug 2024 15:49:24 -0700
Subject: [PATCH] fix more

---
 lib/parse.ml                | 29 +++++++++++++++++++++++------
 lib/websocket_connection.ml | 22 +++++++++++++++++-----
 lib_test/test_httpun_ws.ml  | 10 ++++++----
 3 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/lib/parse.ml b/lib/parse.ml
index 220b9227..8b743c8b 100644
--- a/lib/parse.ml
+++ b/lib/parse.ml
@@ -218,20 +218,37 @@ module Reader = struct
         t.parse_state <- Partial continue
       | _ -> assert false
 
-  let rec read_with_more t bs ~off ~len more =
+  let rec _read_with_more t bs ~off ~len more =
+    let initial = match t.parse_state with Done -> true | _ -> false in
     let consumed =
       match t.parse_state with
       | Fail _ -> 0
+      (* Don't feed empty input when we're at a request boundary *)
+      | Done when len = 0 -> 0
       | Done   ->
         start t (AU.parse t.parser);
-        read_with_more  t bs ~off ~len more;
+        _read_with_more  t bs ~off ~len more;
       | Partial continue ->
         transition t (continue bs more ~off ~len)
     in
-    begin match more with
-    | Complete -> t.closed <- true;
-    | Incomplete -> ()
-    end;
+    (* Special case where the parser just started and was fed a zero-length
+     * bigstring. Avoid putting them parser in an error state in this scenario.
+     * If we were already in a `Partial` state, return the error. *)
+    if initial && len = 0 then t.parse_state <- Done;
+    match t.parse_state with
+    | Done when consumed < len ->
+      let off = off + consumed
+      and len = len - consumed in
+      consumed + _read_with_more t bs ~off ~len more
+    | _ -> consumed
+  ;;
+
+  let read_with_more t bs ~off ~len more =
+    let consumed = _read_with_more t bs ~off ~len more in
+    (match more with
+     | Complete ->
+       t.closed <- true
+     | Incomplete -> ());
     consumed
 
   let force_close t =
diff --git a/lib/websocket_connection.ml b/lib/websocket_connection.ml
index 497fda46..34cc9c69 100644
--- a/lib/websocket_connection.ml
+++ b/lib/websocket_connection.ml
@@ -140,11 +140,26 @@ let next_read_operation t =
 let next_write_operation t =
   Wsd.next t.wsd
 
+let report_exn t exn =
+  set_error_and_handle t (`Exn exn)
+
+let read_with_more t bs ~off ~len more =
+  let consumed = Reader.read_with_more t.reader bs ~off ~len more in
+  if not (Queue.is_empty t.frame_queue)
+  then (
+    let _, payload = Queue.peek t.frame_queue in
+    if Payload.has_pending_output payload
+    then try Payload.execute_read payload
+    with exn -> report_exn t exn
+  );
+  consumed
+;;
+
 let read t bs ~off ~len =
-  Reader.read_with_more t.reader bs ~off ~len Incomplete
+  read_with_more t bs ~off ~len Incomplete
 
 let read_eof t bs ~off ~len =
-  let r = Reader.read_with_more t.reader bs ~off ~len Complete in
+  let r = read_with_more t bs ~off ~len Complete in
   t.eof ();
   r
 
@@ -162,9 +177,6 @@ let yield_writer t k =
 let is_closed { wsd; _ } =
   Wsd.is_closed wsd
 
-let report_exn t exn =
-  set_error_and_handle t (`Exn exn)
-
 let yield_reader t k =
   if Reader.is_closed t.reader
   then k ()
diff --git a/lib_test/test_httpun_ws.ml b/lib_test/test_httpun_ws.ml
index 649de1fa..a73d468f 100644
--- a/lib_test/test_httpun_ws.ml
+++ b/lib_test/test_httpun_ws.ml
@@ -115,10 +115,11 @@ module Websocket = struct
          match opcode with
          | `Text ->
            incr frames_parsed;
-           Payload.schedule_read payload
-             ~on_eof:ignore
-             ~on_read:(fun bs ~off ~len ->
-             Wsd.schedule wsd bs ~kind:`Text ~off ~len)
+           let rec on_read bs ~off ~len =
+             Wsd.schedule wsd ~kind:`Text bs ~off ~len;
+             Payload.schedule_read payload ~on_eof:ignore ~on_read
+           in
+           Payload.schedule_read payload ~on_eof:ignore ~on_read
          | `Binary
          | `Continuation
          | `Connection_close
@@ -144,6 +145,7 @@ module Websocket = struct
      let len = String.length frames in
      let bs = Bigstringaf.of_string ~off:0 ~len frames in
      let read = Server_connection.read t bs ~off:0 ~len in
+     ignore @@ Server_connection.next_read_operation t;
      Alcotest.(check int) "Reads both frames" len read;
      Alcotest.(check int) "Both frames parsed and handled" 2 !frames_parsed;
     ;;