diff --git a/dune b/dune index 00e05f1..dbdc2c7 100644 --- a/dune +++ b/dune @@ -39,6 +39,7 @@ (preprocess (per_module ((pps lwt_ppx) + gzip_io httpev logstash lwt_flag diff --git a/gzip_io.ml b/gzip_io.ml index f962199..55fe25d 100644 --- a/gzip_io.ml +++ b/gzip_io.ml @@ -36,6 +36,31 @@ let string s = IO.nwrite out (Bytes.unsafe_of_string s); (* IO wrong type *) IO.close_out out +let string_lwt ?(chunk_size = 3000) ?(yield = Lwt.pause) s = + let out = output (IO.output_string ()) in + let buff = Buffer.create chunk_size in + let len = String.length s in + let rec loop i = + if i >= len then ( + (* Final flush of the buffer if there's any residue *) + if Buffer.length buff > 0 then IO.nwrite out (Buffer.to_bytes buff); + Lwt.return_unit) + else begin + let c = s.[i] in + Buffer.add_char buff c; + if Buffer.length buff < chunk_size then loop (i + 1) + else ( + (* Buffer is full, write and clear it *) + IO.nwrite out (Buffer.to_bytes buff); + Buffer.clear buff; + (* Yield after processing a chunk *) + let%lwt () = yield () in + loop (i + 1)) + end + in + let%lwt () = loop 0 in + Lwt.return @@ IO.close_out out + let to_string s = let inp = input (IO.input_string s) in let out = IO.output_string () in diff --git a/httpev.ml b/httpev.ml index 1e6b847..5ebcaa8 100644 --- a/httpev.ml +++ b/httpev.ml @@ -928,12 +928,15 @@ let send_reply c cout reply = end in (* possibly apply encoding *) - let (hdrs,body) = + let%lwt (hdrs,body) = (* TODO do not apply encoding to application/gzip *) - (* TODO gzip + chunked? *) match body, code, c.req with - | `Body s, `Ok, Ready { encoding=Gzip; _ } when String.length s > 128 -> ("Content-Encoding", "gzip")::hdrs, `Body (Gzip_io.string s) - | _ -> hdrs, body + | `Body s, `Ok, Ready { encoding=Gzip; _ } when String.length s > 128 -> + let%lwt body = Gzip_io.string_lwt s in + Lwt.return (("Content-Encoding", "gzip")::hdrs, `Body body) + | `Chunks _ as body, `Ok, Ready { encoding=Gzip; _ } -> + Lwt.return (("Content-Encoding", "gzip")::hdrs, body) + | _ -> Lwt.return (hdrs, body) in let hdrs = match body with | `Body s -> ("Content-Length", string_of_int (String.length s)) :: hdrs @@ -955,6 +958,10 @@ let send_reply c cout reply = let push = function | "" -> Lwt.return_unit | s -> + let%lwt s = match c.req with + | Ready { encoding=Gzip; _ } -> Gzip_io.string_lwt s + | _ -> Lwt.return s + in let%lwt () = Lwt_io.write cout (sprintf "%x\r\n" (String.length s)) in let%lwt () = Lwt_io.write cout s in Lwt_io.write cout "\r\n"