Add deflate support for chunked responses

lib/bandit/adapter.ex

@@ -13,6 +13,7 @@ defmodule Bandit.Adapter do
             method: nil,
             status: nil,
             content_encoding: nil,
+            compression_context: nil,
             upgrade: nil,
             metrics: %{},
             opts: []
@@ -24,6 +25,7 @@ defmodule Bandit.Adapter do
           method: Plug.Conn.method() | nil,
           status: Plug.Conn.status() | nil,
           content_encoding: String.t(),
+          compression_context: Bandit.Compression.t() | nil,
           upgrade: nil | {:websocket, opts :: keyword(), websocket_opts :: keyword()},
           metrics: %{},
           opts: %{
@@ -151,7 +153,9 @@ defmodule Bandit.Adapter do
     validate_calling_process!(adapter)
     start_time = Bandit.Telemetry.monotonic_time()
     metrics = Map.put(adapter.metrics, :resp_start_time, start_time)
-    adapter = %{adapter | metrics: metrics}
+
+    {headers, compression_context} = Bandit.Compression.new(adapter, headers, true)
+    adapter = %{adapter | metrics: metrics, compression_context: compression_context}
     {:ok, nil, send_headers(adapter, status, headers, :chunk_encoded)}
   end
 
@@ -168,7 +172,17 @@ defmodule Bandit.Adapter do
     # chunk/2 is unique among Plug.Conn.Adapter's sending callbacks in that it can return an error
     # tuple instead of just raising or dying on error. Rescue here to implement this
     try do
-      {:ok, nil, send_data(adapter, chunk, IO.iodata_length(chunk) == 0)}
+      if IO.iodata_length(chunk) == 0 do
+        compression_metrics = Bandit.Compression.close(adapter.compression_context)
+        adapter = %{adapter | metrics: Map.merge(adapter.metrics, compression_metrics)}
+        {:ok, nil, send_data(adapter, chunk, true)}
+      else
+        {encoded_chunk, compression_context} =
+          Bandit.Compression.compress_chunk(chunk, adapter.compression_context)
+
+        adapter = %{adapter | compression_context: compression_context}
+        {:ok, nil, send_data(adapter, encoded_chunk, false)}
+      end
     rescue
       error -> {:error, Exception.message(error)}
     end

lib/bandit/compression.ex

@@ -20,7 +20,7 @@ defmodule Bandit.Compression do
     |> Enum.find(&(&1 in ~w(deflate gzip x-gzip)))
   end
 
-  def new(adapter, headers) do
+  def new(adapter, headers, streamable \\ false) do
     response_content_encoding_header = Bandit.Headers.get_header(headers, "content-encoding")
 
     response_has_strong_etag =
@@ -45,7 +45,7 @@ defmodule Bandit.Compression do
          !response_has_strong_etag && !response_indicates_no_transform do
       deflate_options = Keyword.get(adapter.opts.http, :deflate_options, [])
 
-      case start_stream(adapter.content_encoding, deflate_options) do
+      case start_stream(adapter.content_encoding, deflate_options, streamable) do
         {:ok, context} -> {[{"content-encoding", adapter.content_encoding} | headers], context}
         {:error, :unsupported_encoding} -> {headers, %__MODULE__{method: :identity}}
       end
@@ -54,7 +54,7 @@ defmodule Bandit.Compression do
     end
   end
 
-  defp start_stream("deflate", opts) do
+  defp start_stream("deflate", opts, _streamable) do
     deflate_context = :zlib.open()
 
     :zlib.deflateInit(
@@ -69,9 +69,9 @@ defmodule Bandit.Compression do
     {:ok, %__MODULE__{method: :deflate, lib_context: deflate_context}}
   end
 
-  defp start_stream("x-gzip", _opts), do: {:ok, %__MODULE__{method: :gzip}}
-  defp start_stream("gzip", _opts), do: {:ok, %__MODULE__{method: :gzip}}
-  defp start_stream(_encoding, _opts), do: {:error, :unsupported_encoding}
+  defp start_stream("x-gzip", _opts, false), do: {:ok, %__MODULE__{method: :gzip}}
+  defp start_stream("gzip", _opts, false), do: {:ok, %__MODULE__{method: :gzip}}
+  defp start_stream(_encoding, _opts, _streamable), do: {:error, :unsupported_encoding}
 
   def compress_chunk(chunk, %__MODULE__{method: :deflate} = context) do
     result = :zlib.deflate(context.lib_context, chunk, :sync)

test/bandit/http1/request_test.exs

@@ -1339,7 +1339,35 @@ defmodule HTTP1RequestTest do
       assert inflated_body == String.duplicate("a", 10_000)
     end
 
+    test "deflate encodes chunk responses", context do
+      response =
+        Req.get!(context.req,
+          url: "/send_big_body_chunked",
+          headers: [{"accept-encoding", "deflate"}]
+        )
+
+      assert response.status == 200
+      assert response.headers["content-encoding"] == ["deflate"]
+      assert response.headers["vary"] == ["accept-encoding"]
 
+      inflate_context = :zlib.open()
+      :ok = :zlib.inflateInit(inflate_context)
+      inflated_body = :zlib.inflate(inflate_context, response.body) |> IO.iodata_to_binary()
+
+      assert inflated_body == String.duplicate("a", 10_000)
+    end
+
+    test "does not gzip encode chunk responses", context do
+      response =
+        Req.get!(context.req,
+          url: "/send_big_body_chunked",
+          headers: [{"accept-encoding", "gzip"}]
+        )
+
+      assert response.status == 200
+      assert response.headers["content-encoding"] == nil
+      assert response.headers["vary"] == ["accept-encoding"]
+      assert response.body == String.duplicate("a", 10_000)
     end
 
     test "falls back to no encoding if no encodings provided", context do
@@ -1444,6 +1472,23 @@ defmodule HTTP1RequestTest do
       |> send_resp(200, String.duplicate("a", 10_000))
     end
 
+    def send_big_body_chunked(conn) do
+      conn = send_chunked(conn, 200)
+
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+      {:ok, conn} = chunk(conn, String.duplicate("a", 1_000))
+
+      conn
+    end
+
     def send_iolist_body(conn) do
       conn
       |> put_resp_header("content-length", "10000")

test/bandit/http2/protocol_test.exs

@@ -809,6 +809,64 @@ defmodule HTTP2ProtocolTest do
       assert SimpleH2Client.recv_body(socket) == {:ok, 1, true, ""}
     end
 
+    test "deflate encodes multiple DATA frames when chunking", context do
+      socket = SimpleH2Client.setup_connection(context)
+
+      headers = [
+        {":method", "GET"},
+        {":path", "/chunk_response"},
+        {":scheme", "https"},
+        {":authority", "localhost:#{context.port}"},
+        {"accept-encoding", "deflate"}
+      ]
+
+      SimpleH2Client.send_headers(socket, 1, true, headers)
+
+      assert {:ok, 1, false,
+              [
+                {":status", "200"},
+                {"date", _date},
+                {"content-encoding", "deflate"},
+                {"vary", "accept-encoding"},
+                {"cache-control", "max-age=0, private, must-revalidate"}
+              ], _ctx} = SimpleH2Client.recv_headers(socket)
+
+      {:ok, 1, false, chunk_1} = SimpleH2Client.recv_body(socket)
+      {:ok, 1, false, chunk_2} = SimpleH2Client.recv_body(socket)
+      assert {:ok, 1, true, ""} == SimpleH2Client.recv_body(socket)
+
+      inflate_context = :zlib.open()
+      :ok = :zlib.inflateInit(inflate_context)
+      inflated_body = :zlib.inflate(inflate_context, [chunk_1, chunk_2]) |> IO.iodata_to_binary()
+
+      assert inflated_body == "OKDOKEE"
+    end
+
+    test "does not gzip encode DATA frames when chunking", context do
+      socket = SimpleH2Client.setup_connection(context)
+
+      headers = [
+        {":method", "GET"},
+        {":path", "/chunk_response"},
+        {":scheme", "https"},
+        {":authority", "localhost:#{context.port}"},
+        {"accept-encoding", "gzip"}
+      ]
+
+      SimpleH2Client.send_headers(socket, 1, true, headers)
+
+      assert {:ok, 1, false,
+              [
+                {":status", "200"},
+                {"date", _date},
+                {"vary", "accept-encoding"},
+                {"cache-control", "max-age=0, private, must-revalidate"}
+              ], _ctx} = SimpleH2Client.recv_headers(socket)
+
+      assert {:ok, 1, false, "OK"} == SimpleH2Client.recv_body(socket)
+      assert {:ok, 1, false, "DOKEE"} == SimpleH2Client.recv_body(socket)
+    end
+
     test "does not write out a body for a chunked response to a HEAD request", context do
       socket = SimpleH2Client.setup_connection(context)
 
@@ -818,6 +876,7 @@ defmodule HTTP2ProtocolTest do
               [
                 {":status", "200"},
                 {"date", _date},
+                {"vary", "accept-encoding"},
                 {"cache-control", "max-age=0, private, must-revalidate"}
               ], _ctx} = SimpleH2Client.recv_headers(socket)
 
@@ -842,6 +901,7 @@ defmodule HTTP2ProtocolTest do
               [
                 {":status", "204"},
                 {"date", _date},
+                {"vary", "accept-encoding"},
                 {"cache-control", "max-age=0, private, must-revalidate"}
               ], _ctx} = SimpleH2Client.recv_headers(socket)
 
@@ -864,6 +924,7 @@ defmodule HTTP2ProtocolTest do
               [
                 {":status", "304"},
                 {"date", _date},
+                {"vary", "accept-encoding"},
                 {"cache-control", "max-age=0, private, must-revalidate"}
               ], _ctx} = SimpleH2Client.recv_headers(socket)