Fix chunked reading for async

lib/transfer_io.ml

@@ -23,36 +23,52 @@ module Make(IO : S.IO) = struct
   type writer = string -> unit IO.t
 
   module Chunked = struct
-    let read ic () =
-      (* Read chunk size *)
-      read_line ic >>= function
-      |Some chunk_size_hex -> begin
-        let chunk_size =
-          let hex =
-            (* chunk size is optionally delimited by ; *)
-            try String.sub chunk_size_hex 0 (String.rindex chunk_size_hex ';')
-            with _ -> chunk_size_hex in
-          try Some (int_of_string ("0x" ^ hex)) with _ -> None
-        in
-        match chunk_size with
+    let read ~remaining ic () =
+      (* read between 0 and 32Kbytes of a chunk *)
+      let read_chunk_fragment () =
+        let max_read_len = Int64.of_int 0x8000 in
+        let len = min !remaining max_read_len in
+        read ic (Int64.to_int len) >>= function
+        | "" -> return ""
+        | buf ->
+          let read_len = Int64.of_int (String.length buf) in
+          remaining := Int64.sub !remaining read_len;
+          (if !remaining = 0L (* End_of_chunk *)
+           then read_line ic (* Junk the CRLF at end of chunk *)
+           else return None) >>= fun _ ->
+          return buf
+      in
+      if !remaining = 0L then
+        (* Beginning of a chunk: read chunk size, read up to 32K bytes *)
+        read_line ic >>= function
         | None -> return Done
-        | Some 0 -> (* TODO: Trailer header support *)
-          let rec read_until_empty_line () =
-            read_line ic >>= function
-            | None | Some "" -> return Done
-            | Some _trailer -> read_until_empty_line ()
-          in
-          read_until_empty_line ()
-        | Some count -> begin
-          read_exactly ic count >>=
-          function
-          |None -> return Done
-          |Some buf ->
-            read_line ic >>= fun _ -> (* Junk the CRLF at end of chunk *)
-            return (Chunk buf)
-        end
-      end
-      |None -> return Done
+        | Some chunk_size_hex -> begin
+            let chunk_size =
+              let hex =
+                (* chunk size is optionally delimited by ; *)
+                try String.sub chunk_size_hex 0 (String.rindex chunk_size_hex ';')
+                with _ -> chunk_size_hex in
+              try Some (Int64.of_string ("0x" ^ hex)) with _ -> None
+            in
+            match chunk_size with
+            | None -> return Done
+            | Some 0L -> (* TODO: Trailer header support *)
+              let rec read_until_empty_line () =
+                read_line ic >>= function
+                | None | Some "" -> return Done
+                | Some _trailer -> read_until_empty_line ()
+              in
+              read_until_empty_line ()
+            | Some count ->
+              remaining := count;
+              read_chunk_fragment () >>= function
+              | "" -> return Done (* 0 bytes read means EOF *)
+              | buf -> return (Chunk buf)
+          end
+      else (* Middle of a chunk, read up to 32K bytes *)
+        read_chunk_fragment () >>= function
+        | "" -> return Done (* 0 bytes read means EOF *)
+        | buf -> return (Chunk buf)
 
     let write oc buf =
       let len = String.length buf in
@@ -100,7 +116,7 @@ module Make(IO : S.IO) = struct
 
   let make_reader =
     function
-    | Chunked -> Chunked.read
+    | Chunked -> Chunked.read ~remaining:(ref 0L)
     | Fixed len -> Fixed.read ~remaining:(ref len)
     | Unknown -> Unknown.read