Improve content-length send logic per RFC9110§8.6/8.7

lib/bandit/headers.ex

@@ -95,33 +95,29 @@ defmodule Bandit.Headers do
         ) ::
           Plug.Conn.headers()
 
-  # For HEAD responses, respect the existing content-length header (if set)
-  # EXCEPT if the status code forbids a content-length header.
-  # If content-length is not set, the server will not send a content-length header.
-  def add_content_length(headers, 0, status, "HEAD") do
-    if send_content_length?(status),
-      do: headers,
-      else: drop_content_length(headers)
+  # Per RFC9110§8.6, we use the following logic:
+  #
+  # * If the response is 1xx or 204, content-length is NEVER sent
+  # * If the response is 304 or the method is HEAD AND the body length is zero, respect any
+  #   content-length header the plug may have set on the assumption that it knows what it would
+  #   have sent
+  # * For all other responses, use the length of the provided response body as the content-length,
+  #   overwriting any content-length the plug may have set
+  def add_content_length(headers, _length, status, _method)
+      when status in 100..199 or status == 204 do
+    drop_content_length(headers)
   end
 
-  # For non-HEAD responses, override the existing content-length header.
-  def add_content_length(headers, length, status, _method) do
-    headers = drop_content_length(headers)
+  def add_content_length(headers, 0, status, method) when status == 304 or method == "HEAD" do
+    headers
+  end
 
-    if send_content_length?(status),
-      do: [{"content-length", to_string(length)} | headers],
-      else: headers
+  def add_content_length(headers, length, _status, _method) do
+    [{"content-length", to_string(length)} | drop_content_length(headers)]
   end
 
   @spec drop_content_length(Plug.Conn.headers()) :: Plug.Conn.headers()
   defp drop_content_length(headers) do
     Enum.reject(headers, &(elem(&1, 0) == "content-length"))
   end
-
-  # Per RFC9110§8.6
-  @spec send_content_length?(Plug.Conn.int_status()) :: boolean()
-  defp send_content_length?(status) when status in 100..199, do: false
-  defp send_content_length?(204), do: false
-  defp send_content_length?(304), do: false
-  defp send_content_length?(_), do: true
 end

test/bandit/http1/request_test.exs

@@ -1542,17 +1542,45 @@ defmodule HTTP1RequestTest do
       send_resp(conn, 204, "this is an invalid body")
     end
 
-    test "writes out a response with no content-length header or body for 304 responses",
+    test "writes out a response with content-length header but no body for 304 responses",
          context do
       client = SimpleHTTP1Client.tcp_client(context)
       SimpleHTTP1Client.send(client, "GET", "/send_304", ["host: localhost"])
 
-      assert {:ok, "304 Not Modified", headers, ""} = SimpleHTTP1Client.recv_reply(client)
-      assert Bandit.Headers.get_header(headers, :"content-length") == nil
+      assert {:ok, "304 Not Modified", headers, ""} = SimpleHTTP1Client.recv_reply(client, true)
+      assert Bandit.Headers.get_header(headers, :"content-length") == "5"
     end
 
     def send_304(conn) do
-      send_resp(conn, 304, "this is an invalid body")
+      send_resp(conn, 304, "abcde")
+    end
+
+    test "respects plug-provided zero content-length and no body for 304 responses", context do
+      response = Req.head!(context.req, url: "/send_304_zero_content_length")
+
+      assert response.status == 304
+      assert response.body == ""
+      assert response.headers["content-length"] == ["0"]
+    end
+
+    def send_304_zero_content_length(conn) do
+      conn
+      |> put_resp_header("content-length", "0")
+      |> send_resp(304, "")
+    end
+
+    test "respects plug-provided nonzero content-length but no body for 304 responses", context do
+      response = Req.head!(context.req, url: "/send_304_nonzero_content_length")
+
+      assert response.status == 304
+      assert response.body == ""
+      assert response.headers["content-length"] == ["5"]
+    end
+
+    def send_304_nonzero_content_length(conn) do
+      conn
+      |> put_resp_header("content-length", "5")
+      |> send_resp(304, "abcde")
     end
 
     test "writes out a response with zero content-length for 200 responses", context do
@@ -1575,7 +1603,8 @@ defmodule HTTP1RequestTest do
       send_resp(conn, 200, "")
     end
 
-    test "writes out a response with zero content-length for HEAD 200 responses", context do
+    test "respects plug-provided zero content-length and no body for HEAD 200 responses",
+         context do
       response = Req.head!(context.req, url: "/send_200_zero_content_length")
 
       assert response.status == 200
@@ -1589,6 +1618,21 @@ defmodule HTTP1RequestTest do
       |> send_resp(200, "")
     end
 
+    test "respects plug-provided nonzero content-length but no body for HEAD 200 responses",
+         context do
+      response = Req.head!(context.req, url: "/send_200_nonzero_content_length")
+
+      assert response.status == 200
+      assert response.body == ""
+      assert response.headers["content-length"] == ["5"]
+    end
+
+    def send_200_nonzero_content_length(conn) do
+      conn
+      |> put_resp_header("content-length", "5")
+      |> send_resp(200, "abcde")
+    end
+
     test "writes out a response with zero content-length for 301 responses", context do
       response = Req.get!(context.req, url: "/send_301")
 
@@ -1712,13 +1756,12 @@ defmodule HTTP1RequestTest do
       |> send_file(204, Path.join([__DIR__, "../../support/sendfile"]), 0, :all)
     end
 
-    test "does not write out a content-length header or body for files on a 304",
-         context do
+    test "write out a content-length header but no body for files on a 304", context do
       client = SimpleHTTP1Client.tcp_client(context)
       SimpleHTTP1Client.send(client, "HEAD", "/send_full_file_304", ["host: localhost"])
 
       assert {:ok, "304 Not Modified", headers, ""} = SimpleHTTP1Client.recv_reply(client, true)
-      assert Bandit.Headers.get_header(headers, :"content-length") == nil
+      assert Bandit.Headers.get_header(headers, :"content-length") == "6"
       assert SimpleHTTP1Client.connection_closed_for_reading?(client)
     end
 

test/bandit/http2/protocol_test.exs

@@ -723,7 +723,7 @@ defmodule HTTP2ProtocolTest do
       send_resp(conn, 204, "this is an invalid body")
     end
 
-    test "sends no content-length header or body for 304 responses", context do
+    test "sends content-length header but no body for 304 responses", context do
       socket = SimpleH2Client.setup_connection(context)
 
       SimpleH2Client.send_simple_headers(socket, 1, :get, "/send_304", context[:port])
@@ -732,13 +732,14 @@ defmodule HTTP2ProtocolTest do
               [
                 {":status", "304"},
                 {"date", _date},
+                {"content-length", "5"},
                 {"vary", "accept-encoding"},
                 {"cache-control", "max-age=0, private, must-revalidate"}
               ], _ctx} = SimpleH2Client.recv_headers(socket)
     end
 
     def send_304(conn) do
-      send_resp(conn, 304, "this is an invalid body")
+      send_resp(conn, 304, "abcde")
     end
 
     test "sends headers but no body for a HEAD request to a file", context do
@@ -791,6 +792,7 @@ defmodule HTTP2ProtocolTest do
               [
                 {":status", "304"},
                 {"date", _date},
+                {"content-length", "6"},
                 {"cache-control", "max-age=0, private, must-revalidate"}
               ], _ctx} = SimpleH2Client.recv_headers(socket)