From e6b46379c0a678df00a97f65b2a459b9d302040b Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sun, 17 Nov 2019 23:55:42 -0600
Subject: [PATCH] wip(zlib): implement decode/encode wrappers

---
 src/zlib/Tiny_httpd_zlib.ml | 88 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 86 insertions(+), 2 deletions(-)

diff --git a/src/zlib/Tiny_httpd_zlib.ml b/src/zlib/Tiny_httpd_zlib.ml
index bc59a1fd..8dcbadb8 100644
--- a/src/zlib/Tiny_httpd_zlib.ml
+++ b/src/zlib/Tiny_httpd_zlib.ml
@@ -5,8 +5,8 @@ module BS = Tiny_httpd.Byte_stream
 let mk_decode_deflate_stream_ () (is:S.byte_stream) : S.byte_stream =
   S._debug (fun k->k "wrap stream with inflate.decode");
   let buf = Bytes.make 4096 ' ' in
-  let write_offset = ref 0 in
   let buf_len = ref 0 in
+  let write_offset = ref 0 in
   let zlib_str = Zlib.inflate_init false in
   let is_done = ref false in
   let bs_close () =
@@ -49,6 +49,74 @@ let mk_decode_deflate_stream_ () (is:S.byte_stream) : S.byte_stream =
   in
   {S.bs_fill_buf; bs_consume; bs_close}
 
+let encode_deflate_stream_ (is:S.byte_stream) : S.byte_stream =
+  S._debug (fun k->k "wrap stream with deflate.encode");
+  let refill = ref true in
+  let buf = Bytes.make 4096 ' ' in
+  let buf_len = ref 0 in
+  let write_offset = ref 0 in
+  let zlib_str = Zlib.deflate_init 4 false in
+  let bs_close () =
+    Zlib.deflate_end zlib_str;
+    BS.close is
+  in
+  let bs_consume n =
+    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 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 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
+  in
+  {S.bs_fill_buf; bs_consume; bs_close}
+
+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
@@ -72,9 +140,25 @@ let cb_decode_compressed_stream (req:unit S.Request.t) : _ option =
     end
   | _ -> None
 
+let cb_encode_compressed_stream (req:_ S.Request.t) (resp:S.Response.t) : _ option =
+  if accept_deflate req then (
+    match resp.body with
+    | `String _ -> None
+    | `Stream str ->
+      S._debug (fun k->k "encode response with deflate");
+      Some {
+        resp with
+        headers=
+          (resp.headers
+           |> S.Headers.remove "Content-Length"
+           |> S.Headers.set "Transfer-Coding" "deflate; chunked");
+        body=`Stream (encode_deflate_stream_ str);
+      }
+  ) else None
+
 let setup (server:S.t) : unit =
   S._debug (fun k->k "setup gzip support");
   S.add_decode_request_cb server cb_decode_compressed_stream;
-  (*   S.add_encode_response_cb server (assert false); *)
+  S.add_encode_response_cb server cb_encode_compressed_stream;
   ()