diff --git a/src/Tiny_httpd.ml b/src/Tiny_httpd.ml index 73dadceb..84a0d6fb 100644 --- a/src/Tiny_httpd.ml +++ b/src/Tiny_httpd.ml @@ -102,7 +102,7 @@ module Byte_stream = struct ref ( match len with | Some n -> - if n > Bytes.length s -i then invalid_arg "Byte_stream.of_bytes"; + if n > Bytes.length s - i then invalid_arg "Byte_stream.of_bytes"; n | None -> Bytes.length s - i ) @@ -110,7 +110,7 @@ module Byte_stream = struct let i = ref i in { bs_fill_buf=(fun () -> s, !i, !len); bs_close=(fun () -> len := 0); - bs_consume=(fun n -> assert (n<= !len); i := !i + n; len := !len - n); + bs_consume=(fun n -> assert (n>=0 && n<= !len); i := !i + n; len := !len - n); } let of_string s : t = diff --git a/src/camlzip/Tiny_httpd_camlzip.ml b/src/camlzip/Tiny_httpd_camlzip.ml index d6659b26..28c6ff31 100644 --- a/src/camlzip/Tiny_httpd_camlzip.ml +++ b/src/camlzip/Tiny_httpd_camlzip.ml @@ -64,38 +64,59 @@ let encode_deflate_stream_ ~buf_size (is:S.byte_stream) : S.byte_stream = write_offset := n + !write_offset in let bs_fill_buf () = - S._debug (fun k->k "deflate.fill buf"); - if !write_offset >= !buf_len then ( - (* see if we need to refill *) - write_offset := 0; - buf_len := 0; - if !refill then ( + let rec loop() = + S._debug (fun k->k "deflate.fill.iter out_off=%d out_len=%d" + !write_offset !buf_len); + if !write_offset < !buf_len then ( + (* still the same slice, not consumed entirely by output *) + buf, !write_offset, !buf_len - !write_offset + ) else if not !refill then ( + (* empty slice, no refill *) + buf, !write_offset, !buf_len - !write_offset + ) else ( + (* the output was entirely consumed, we need to do more work *) + write_offset := 0; + buf_len := 0; let in_s, in_i, in_len = is.S.bs_fill_buf () in - begin - try - (* decompress from input buffer *) - let finished, used_in, used_out = - Zlib.deflate zlib_str - in_s in_i in_len - buf 0 (Bytes.length buf) - Zlib.Z_SYNC_FLUSH - in - buf_len := used_out; - is.S.bs_consume used_in; - S._debug - (fun k->k "encode %d bytes as %d bytes using deflate (finished: %b)" - used_in used_out finished); - if finished then ( - S._debug (fun k->k "deflate: finished"); - refill := false; - ); - with Zlib.Error (e1,e2) -> - S.Response.fail_raise ~code:400 - "deflate: error during compression:\n%s %s" e1 e2 - end; - ); - ); - buf, !write_offset, !buf_len - !write_offset + if in_len>0 then ( + (* try to decompress from input buffer *) + let _finished, used_in, used_out = + Zlib.deflate zlib_str + in_s in_i in_len + buf 0 (Bytes.length buf) + Zlib.Z_NO_FLUSH + in + buf_len := used_out; + is.S.bs_consume used_in; + S._debug + (fun k->k "encode %d bytes as %d bytes using deflate (finished: %b)" + used_in used_out _finished); + if _finished then ( + S._debug (fun k->k "deflate: finished"); + refill := false; + ); + loop() + ) else ( + (* finish sending the internal state *) + let _finished, used_in, used_out = + Zlib.deflate zlib_str + in_s in_i in_len + buf 0 (Bytes.length buf) + Zlib.Z_FULL_FLUSH + in + assert (used_in = 0); + buf_len := used_out; + if used_out = 0 then ( + refill := false; + ); + loop() + ) + ) + in + try loop() + with Zlib.Error (e1,e2) -> + S.Response.fail_raise ~code:400 + "deflate: error during compression:\n%s %s" e1 e2 in {S.bs_fill_buf; bs_consume; bs_close} @@ -141,22 +162,25 @@ let cb_decode_compressed_stream ~buf_size (req:unit S.Request.t) : _ option = | _ -> None let cb_encode_compressed_stream + ~compress_above ~buf_size (req:_ S.Request.t) (resp:S.Response.t) : _ option = if accept_deflate req then ( let set_headers h = h |> S.Headers.remove "Content-Length" - |> S.Headers.set "Content-Encoding" "deflate" + |> S.Headers.set "Content-Encoding" "deflate, chunked" in match resp.body with - | `String s when String.length s > 500 * 1024 -> - S._debug (fun k->k "encode str response with deflate"); + | `String s when String.length s > compress_above -> + S._debug + (fun k->k "encode str response with deflate (size %d, threshold %d)" + (String.length s) compress_above); let body = encode_deflate_stream_ ~buf_size @@ S.Byte_stream.of_string s in Some { resp with - headers= set_headers resp.headers; body=`Stream body; + headers=set_headers resp.headers; body=`Stream body; } | `Stream str -> S._debug (fun k->k "encode stream response with deflate"); @@ -168,10 +192,12 @@ let cb_encode_compressed_stream | `String _ -> None ) else None -let setup ?(buf_size=48 * 1_024) (server:S.t) : unit = +let setup + ?(compress_above=500*1024) + ?(buf_size=48 * 1_024) (server:S.t) : unit = let buf_size = max buf_size 1_024 in S._debug (fun k->k "setup gzip support (buf-size %d)" buf_size); S.add_decode_request_cb server (cb_decode_compressed_stream ~buf_size); - S.add_encode_response_cb server (cb_encode_compressed_stream ~buf_size); + S.add_encode_response_cb server (cb_encode_compressed_stream ~compress_above ~buf_size); () diff --git a/src/camlzip/Tiny_httpd_camlzip.mli b/src/camlzip/Tiny_httpd_camlzip.mli index fda04b07..dd2e3cb6 100644 --- a/src/camlzip/Tiny_httpd_camlzip.mli +++ b/src/camlzip/Tiny_httpd_camlzip.mli @@ -1,5 +1,8 @@ -val setup : ?buf_size:int -> Tiny_httpd.t -> unit +val setup : + ?compress_above:int -> + ?buf_size:int -> Tiny_httpd.t -> unit (** Install callbacks for tiny_httpd to be able to encode/decode compressed streams + @param compress_above threshold above with string responses are compressed @param buf_size size of the underlying buffer for compression/decompression *)