transparent (content-encoding) compression

benoitc · Nov 28, 2017 · 06e6cb3 · 06e6cb3
1 parent afbb7e5
commit 06e6cb3
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 15 deletions.
diff --git a/include/hackney_lib.hrl b/include/hackney_lib.hrl
@@ -26,8 +26,10 @@
  partial_headers = []   :: list(),
  clen                   :: integer() | undefined,
  te = <<>>              :: binary(),
+ ce = <<>>              :: binary(),
  connection = <<>>      :: binary(),
  ctype = <<>>           :: binary(),
  location = <<>>        :: binary(),
- body_state = waiting   :: atom() | tuple()
+ body_state = waiting   :: atom() | tuple(),
+ encoding               :: {atom(), any()} | undefined
 }).
diff --git a/src/hackney.erl b/src/hackney.erl
@@ -92,6 +92,7 @@ cancel_request(Ref) ->
 
 %% @doc set client options.
 %% Options are:
+%% - `compress': request compression and transparently decompress
 %% - `async': to fetch the response asynchronously
 %% - `{async, once}': to receive the response asynchronously once time.
 %% To receive the next message use the function `hackney:stream_next/1'.
@@ -200,6 +201,8 @@ request(Method, URL, Headers, Body) ->
 %%      directly. The response is `{ok, Status, Headers, Body}'</li>
 %%      <li>`max_body': sets maximum allowed size of the body if
 %%      with_body is true</li>
+%%      <li>`compress': request that the server sends the body compressed
+%%      and instructs hackney to transparently decompress it</li>
 %%      <li>`async': receive the response asynchronously
 %%      The function return {ok, StreamRef}.
 %%      When {async, once} is used the response will be received only once. To
@@ -968,6 +971,9 @@ maybe_update_req(State) ->
 
 parse_options([], State) ->
   State;
+parse_options([compress | Rest], State = #client{headers=Headers}) ->
+  Headers2 = hackney_headers:store(<<"accept-encoding">>, <<"gzip, deflate">>, Headers),
+  parse_options(Rest, State#client{headers=Headers2});
 parse_options([async | Rest], State) ->
   parse_options(Rest, State#client{async=true});
 parse_options([{async, Async} | Rest], State) ->

diff --git a/src/hackney_http.erl b/src/hackney_http.erl
@@ -300,6 +300,9 @@ parse_header(Line, St) ->
           <<"transfer-encoding">> ->
             TE = hackney_bstr:to_lower(hackney_bstr:trim(Value)),
             St#hparser{te=TE};
+          <<"content-encoding">> ->
+            CE = hackney_bstr:to_lower(hackney_bstr:trim(Value)),
+            St#hparser{ce=CE};
           <<"connection">> ->
             Connection = hackney_bstr:to_lower(hackney_bstr:trim(Value)),
             St#hparser{connection=Connection};
@@ -325,7 +328,40 @@ parse_trailers(St, Acc) ->
     _ -> error
   end.
 
-parse_body(St=#hparser{body_state=waiting, te=TE, clen=Length,
+-define(MAX_WBITS, 15).	% zconf.h
+parse_body(St=#hparser{body_state=waiting, ce=CE, clen=Length, method=Method}) ->
+  St2 = case CE of
+    _ when Length =:= 0 orelse Method =:= <<"HEAD">> ->
+      St;
+    <<"gzip">> ->
+      Z = zlib:open(),
+      ok = zlib:inflateInit(Z, 16 + ?MAX_WBITS),	% http://www.zlib.net/manual.html#Advanced
+      St#hparser{encoding={zlib,Z}};
+    <<"deflate">> ->
+      Z = zlib:open(),
+      ok = zlib:inflateInit(Z, ?MAX_WBITS),
+      St#hparser{encoding={zlib,Z}};
+    _ ->
+      St
+  end,
+  parse_body2(St2);
+parse_body(St=#hparser{encoding={zlib,Z}}) ->
+  case parse_body2(St) of
+    {ok, Chunk, St2} ->
+      Chunk2 = iolist_to_binary(zlib:inflate(Z, Chunk)),
+      {ok, Chunk2, St2};
+    {done, Rest} ->
+      Rest2 = iolist_to_binary(zlib:inflate(Z, Rest)),
+      ok = zlib:inflateEnd(Z),
+      ok = zlib:close(Z),
+      {done, Rest2};
+    Else ->
+      Else
+  end;
+parse_body(St) ->
+  parse_body2(St).
+
+parse_body2(St=#hparser{body_state=waiting, te=TE, clen=Length,
   method=Method, buffer=Buffer}) ->
   case TE of
     <<"chunked">> ->
@@ -340,15 +376,14 @@ parse_body(St=#hparser{body_state=waiting, te=TE, clen=Length,
     _ ->
       {done, Buffer}
   end;
-parse_body(#hparser{body_state=done, buffer=Buffer}) ->
+parse_body2(#hparser{body_state=done, buffer=Buffer}) ->
   {done, Buffer};
-parse_body(St=#hparser{buffer=Buffer, body_state={stream, _, _, _}})
+parse_body2(St=#hparser{buffer=Buffer, body_state={stream, _, _, _}})
   when byte_size(Buffer) > 0 ->
   transfer_decode(Buffer, St#hparser{buffer= <<>>});
-parse_body(St) ->
+parse_body2(St) ->
   {more, St, <<>>}.
 
-
 -spec transfer_decode(binary(), #hparser{})
     -> {ok, binary(), #hparser{}} | {error, atom()}.
 transfer_decode(Data, St=#hparser{
@@ -510,6 +545,8 @@ get_property(method, #hparser{method=Method}) ->
   Method;
 get_property(transfer_encoding, #hparser{te=TE}) ->
   TE;
+get_property(content_encoding, #hparser{ce=CE}) ->
+  CE;
 get_property(content_length, #hparser{clen=CLen}) ->
   CLen;
 get_property(connection, #hparser{connection=Connection}) ->

diff --git a/src/hackney_request.erl b/src/hackney_request.erl
@@ -57,28 +57,33 @@ perform(Client0, {Method0, Path, Headers0, Body0}) ->
   %% add host eventually
   Headers2 = maybe_add_host(Headers1, Client0#client.netloc),
 
+  Compress = proplists:get_value(compress, Options, false),
+  Headers3 = if Compress ->
+    hackney_headers_new:store(<<"accept-encoding">>, <<"gzip, deflate">>, Headers2);
+      true -> Headers2 end,
+
   %% get expect headers
-  Expect = expectation(Headers2),
+  Expect = expectation(Headers3),
 
   %% build headers with the body.
   {FinalHeaders, ReqType, Body, Client1} = case Body0 of
                                             stream ->
-                                               {Headers2, ReqType0, stream, Client0};
+                                               {Headers3, ReqType0, stream, Client0};
                                             stream_multipart ->
-                                              handle_multipart_body(Headers2, ReqType0, Client0);
+                                              handle_multipart_body(Headers3, ReqType0, Client0);
                                             {stream_multipart, Size} ->
-                                              handle_multipart_body(Headers2, ReqType0, Size, Client0);
+                                              handle_multipart_body(Headers3, ReqType0, Size, Client0);
                                             {stream_multipart, Size, Boundary} ->
-                                              handle_multipart_body(Headers2, ReqType0,
+                                              handle_multipart_body(Headers3, ReqType0,
                                                                     Size, Boundary, Client0);
                                             <<>> when Method =:= <<"POST">> orelse Method =:= <<"PUT">> ->
-                                              handle_body(Headers2, ReqType0, Body0, Client0);
+                                              handle_body(Headers3, ReqType0, Body0, Client0);
                                             <<>> ->
-                                              {Headers2, ReqType0, Body0, Client0};
+                                              {Headers3, ReqType0, Body0, Client0};
                                             [] ->
-                                              {Headers2, ReqType0, Body0, Client0};
+                                              {Headers3, ReqType0, Body0, Client0};
                                             _ ->
-                                              handle_body(Headers2, ReqType0, Body0, Client0)
+                                              handle_body(Headers3, ReqType0, Body0, Client0)
                                           end,
 
   %% build final client record