diff --git a/async/cohttp_async_io.ml b/async/cohttp_async_io.ml index 6490545d6c..1e6deeae00 100644 --- a/async/cohttp_async_io.ml +++ b/async/cohttp_async_io.ml @@ -50,12 +50,16 @@ let read_line = ) let read ic len = + if len > 0x8000 then + invalid_arg "read: trying to allocate more than 0x8000 bytes"; let buf = String.create len in Reader.read ic ~len buf >>| function | `Ok len' -> String.sub buf 0 len' | `Eof -> "" let read_exactly ic len = + if len > 0x8000 then + invalid_arg "read: trying to allocate more than 0x8000 bytes"; let buf = String.create len in Reader.really_read ic ~pos:0 ~len buf >>| function diff --git a/lib/transfer_io.ml b/lib/transfer_io.ml index eb8651dc2f..e1b4b5be39 100644 --- a/lib/transfer_io.ml +++ b/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 then (* End_of_chunk *) + 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 + | 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 | 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 + 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