tiny_httpd/src/camlzip/Tiny_httpd_camlzip.ml

module S = Tiny_httpd_server
module BS = Tiny_httpd_stream

let decode_deflate_stream_ ~buf_size (is : S.byte_stream) : S.byte_stream =
  S._debug (fun k -> k "wrap stream with deflate.decode");
  let zlib_str = Zlib.inflate_init false in
  let is_done = ref false in
  BS.make ~bs:(Bytes.create buf_size)
    ~close:(fun _ ->
      Zlib.inflate_end zlib_str;
      BS.close is)
    ~consume:(fun self len ->
      if len > self.len then
        S.Response.fail_raise ~code:400
          "inflate: error during decompression: invalid consume len %d (max %d)"
          len self.len;
      self.off <- self.off + len;
      self.len <- self.len - len)
    ~fill:(fun self ->
      (* refill [buf] if needed *)
      if self.len = 0 && not !is_done then (
        is.fill_buf ();
        (try
           let finished, used_in, used_out =
             Zlib.inflate zlib_str self.bs 0 (Bytes.length self.bs) is.bs is.off
               is.len Zlib.Z_SYNC_FLUSH
           in
           is.consume used_in;
           self.off <- 0;
           self.len <- used_out;
           if finished then is_done := true;
           S._debug (fun k ->
               k "decode %d bytes as %d bytes from inflate (finished: %b)"
                 used_in used_out finished)
         with Zlib.Error (e1, e2) ->
           S.Response.fail_raise ~code:400
             "inflate: error during decompression:\n%s %s" e1 e2);
        S._debug (fun k ->
            k "inflate: refill %d bytes into internal buf" self.len)
      ))
    ()

let encode_deflate_stream_ ~buf_size (is : S.byte_stream) : S.byte_stream =
  S._debug (fun k -> k "wrap stream with deflate.encode");
  let refill = ref true in
  let zlib_str = Zlib.deflate_init 4 false in
  BS.make ~bs:(Bytes.create buf_size)
    ~close:(fun _self ->
      S._debug (fun k -> k "deflate: close");
      Zlib.deflate_end zlib_str;
      BS.close is)
    ~consume:(fun self n ->
      self.off <- self.off + n;
      self.len <- self.len - n)
    ~fill:(fun self ->
      let rec loop () =
        S._debug (fun k ->
            k "deflate.fill.iter out_off=%d out_len=%d" self.off self.len);
        if self.len > 0 then
          ()
        (* still the same slice, not consumed entirely by output *)
        else if not !refill then
          ()
        (* empty slice, no refill *)
        else (
          (* the output was entirely consumed, we need to do more work *)
          is.BS.fill_buf ();
          if is.len > 0 then (
            (* try to decompress from input buffer *)
            let _finished, used_in, used_out =
              Zlib.deflate zlib_str is.bs is.off is.len self.bs 0
                (Bytes.length self.bs) Zlib.Z_NO_FLUSH
            in
            self.off <- 0;
            self.len <- used_out;
            is.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 (
            (* [is] is done, finish sending the data in current buffer *)
            let _finished, used_in, used_out =
              Zlib.deflate zlib_str is.bs is.off is.len self.bs 0
                (Bytes.length self.bs) Zlib.Z_FULL_FLUSH
            in
            assert (used_in = 0);
            self.off <- 0;
            self.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)
    ()

let split_on_char ?(f = fun x -> x) c s : string list =
  let rec loop acc i =
    match String.index_from s i c with
    | exception Not_found ->
      let acc =
        if i = String.length s then
          acc
        else
          f (String.sub s i (String.length s - i)) :: acc
      in
      List.rev acc
    | j ->
      let acc = f (String.sub s i (j - i)) :: acc in
      loop acc (j + 1)
  in
  loop [] 0

let accept_deflate (req : _ S.Request.t) =
  match S.Request.get_header req "Accept-Encoding" with
  | Some s -> List.mem "deflate" @@ split_on_char ~f:String.trim ',' s
  | None -> false

let has_deflate s =
  try Scanf.sscanf s "deflate, %s" (fun _ -> true) with _ -> false

(* decompress [req]'s body if needed *)
let decompress_req_stream_ ~buf_size (req : BS.t S.Request.t) : _ S.Request.t =
  match S.Request.get_header ~f:String.trim req "Transfer-Encoding" with
  (* TODO
     | Some "gzip" ->
       let req' = S.Request.set_header req "Transfer-Encoding" "chunked" in
       Some (req', decode_gzip_stream_)
  *)
  | Some s when has_deflate s ->
    (match Scanf.sscanf s "deflate, %s" (fun s -> s) with
    | tr' ->
      let body' = S.Request.body req |> decode_deflate_stream_ ~buf_size in
      req
      |> S.Request.set_header "Transfer-Encoding" tr'
      |> S.Request.set_body body'
    | exception _ -> req)
  | _ -> req

let compress_resp_stream_ ~compress_above ~buf_size (req : _ S.Request.t)
    (resp : S.Response.t) : S.Response.t =
  (* headers for compressed stream *)
  let update_headers h =
    h
    |> S.Headers.remove "Content-Length"
    |> S.Headers.set "Content-Encoding" "deflate"
  in

  if accept_deflate req then (
    match resp.body with
    | `String s when String.length s > compress_above ->
      (* big string, we compress *)
      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 @@ BS.of_string s in
      resp
      |> S.Response.update_headers update_headers
      |> S.Response.set_body (`Stream body)
    | `Stream str ->
      S._debug (fun k -> k "encode stream response with deflate");
      resp
      |> S.Response.update_headers update_headers
      |> S.Response.set_body (`Stream (encode_deflate_stream_ ~buf_size str))
    | `String _ | `Void -> resp
  ) else
    resp

let middleware ?(compress_above = 16 * 1024) ?(buf_size = 16 * 1_024) () :
    S.Middleware.t =
  let buf_size = max buf_size 1_024 in
  fun h req ~resp ->
    let req = decompress_req_stream_ ~buf_size req in
    h req ~resp:(fun response ->
        resp @@ compress_resp_stream_ ~buf_size ~compress_above req response)

let setup ?compress_above ?buf_size server =
  let m = middleware ?compress_above ?buf_size () in
  S._debug (fun k -> k "setup gzip support");
  S.add_middleware ~stage:`Encoding server m