Implement max_concurrent_streams (#524)

* Implement max_concurrent_streams literally

* Ensure max_requests=0 means no limit

* Raise earlier and avoid a few conditionals for performance

Author: NelsonVides

lib/bandit/http2/connection.ex

@@ -222,42 +222,51 @@ defmodule Bandit.HTTP2.Connection do
         connection.streams
 
       :new ->
-        if accept_stream?(connection) do
-          stream =
-            Bandit.HTTP2.Stream.init(
-              self(),
-              stream_id,
-              connection.remote_settings.initial_window_size
-            )
-
-          case Bandit.HTTP2.StreamProcess.start_link(
-                 stream,
-                 connection.plug,
-                 connection.telemetry_span,
-                 connection.conn_data,
-                 connection.opts
-               ) do
-            {:ok, pid} ->
-              streams = Bandit.HTTP2.StreamCollection.insert(connection.streams, stream_id, pid)
-              with_stream(%{connection | streams: streams}, stream_id, fun)
-
-            _ ->
-              raise "Unable to start stream process"
-          end
-        else
-          connection_error!("Connection count exceeded", Bandit.HTTP2.Errors.refused_stream())
+        new_stream!(connection, stream_id)
+
+        stream =
+          Bandit.HTTP2.Stream.init(
+            self(),
+            stream_id,
+            connection.remote_settings.initial_window_size
+          )
+
+        case Bandit.HTTP2.StreamProcess.start_link(
+               stream,
+               connection.plug,
+               connection.telemetry_span,
+               connection.conn_data,
+               connection.opts
+             ) do
+          {:ok, pid} ->
+            streams = Bandit.HTTP2.StreamCollection.insert(connection.streams, stream_id, pid)
+            with_stream(%{connection | streams: streams}, stream_id, fun)
+
+          _ ->
+            raise "Unable to start stream process"
         end
 
       :invalid ->
         connection_error!("Received invalid stream identifier")
     end
   end
 
-  defp accept_stream?(connection) do
+  defp new_stream!(connection, stream_id) do
     max_requests = Keyword.get(connection.opts.http_2, :max_requests, 0)
 
-    max_requests == 0 ||
-      Bandit.HTTP2.StreamCollection.stream_count(connection.streams) < max_requests
+    if max_requests != 0 and
+         max_requests <= Bandit.HTTP2.StreamCollection.stream_count(connection.streams) do
+      connection_error!("Connection count exceeded", Bandit.HTTP2.Errors.refused_stream())
+    end
+
+    if connection.local_settings.max_concurrent_streams <=
+         Bandit.HTTP2.StreamCollection.open_stream_count(connection.streams) do
+      stream_error!(
+        "Concurrent stream count exceeded",
+        stream_id,
+        Bandit.HTTP2.Errors.refused_stream()
+      )
+    end
   end
 
   defp check_oversize_fragment!(fragment, connection) do
@@ -401,6 +410,19 @@ defmodule Bandit.HTTP2.Connection do
     raise Bandit.HTTP2.Errors.ConnectionError, message: message, error_code: error_code
   end
 
+  @spec stream_error!(
+          String.t(),
+          Bandit.HTTP2.Stream.stream_id(),
+          Bandit.HTTP2.Errors.error_code()
+        ) ::
+          no_return()
+  defp stream_error!(message, stream_id, error_code) do
+    raise Bandit.HTTP2.Errors.StreamError,
+      message: message,
+      error_code: error_code,
+      stream_id: stream_id
+  end
+
   defp send_frame(frame, socket, connection) do
     _ =
       ThousandIsland.Socket.send(

lib/bandit/http2/handler.ex

@@ -14,7 +14,7 @@ defmodule Bandit.HTTP2.Handler do
     connection = Bandit.HTTP2.Connection.init(socket, state.plug, state.opts)
     {:continue, Map.merge(state, %{buffer: <<>>, connection: connection})}
   rescue
-    error -> rescue_error(error, __STACKTRACE__, socket, state)
+    error -> rescue_connection_error(error, __STACKTRACE__, socket, state)
   end
 
   @impl ThousandIsland.Handler
@@ -38,7 +38,8 @@ defmodule Bandit.HTTP2.Handler do
     end)
     |> then(&{:continue, &1})
   rescue
-    error -> rescue_error(error, __STACKTRACE__, socket, state)
+    error in Bandit.HTTP2.Errors.StreamError -> rescue_stream_error(error, socket, state)
+    error -> rescue_connection_error(error, __STACKTRACE__, socket, state)
   end
 
   @impl ThousandIsland.Handler
@@ -148,7 +149,18 @@ defmodule Bandit.HTTP2.Handler do
     {:noreply, {socket, %{state | connection: connection}}, socket.read_timeout}
   end
 
-  defp rescue_error(error, stacktrace, socket, state) do
+  defp rescue_stream_error(error, socket, state) do
+    Bandit.HTTP2.Connection.send_rst_stream(
+      error.stream_id,
+      error.error_code,
+      socket,
+      state.connection
+    )
+
+    {:continue, state}
+  end
+
+  defp rescue_connection_error(error, stacktrace, socket, state) do
     do_rescue_error(error, stacktrace, socket, state)
     {:close, state}
   end

lib/bandit/http2/stream_collection.ex

@@ -70,6 +70,9 @@ defmodule Bandit.HTTP2.StreamCollection do
   @spec stream_count(t()) :: non_neg_integer()
   def stream_count(collection), do: collection.stream_count
 
+  @spec open_stream_count(t()) :: non_neg_integer()
+  def open_stream_count(collection), do: collection.pid_to_id |> map_size()
+
   @spec last_stream_id(t()) :: Bandit.HTTP2.Stream.stream_id()
   def last_stream_id(collection), do: collection.last_stream_id
 end

test/bandit/http2/protocol_test.exs

@@ -1,6 +1,7 @@
 defmodule HTTP2ProtocolTest do
   use ExUnit.Case, async: true
   use ServerHelpers
+  use Machete
 
   import Bitwise
 
@@ -233,6 +234,16 @@ defmodule HTTP2ProtocolTest do
       assert_receive {:log, %{level: :error, msg: {:string, msg}}}, 500
       assert msg == "** (Bandit.HTTP2.Errors.ConnectionError) Connection count exceeded"
     end
+
+    @tag :capture_log
+    test "max_requests zero does not put a limit", context do
+      context = https_server(context, http_2_options: [max_requests: 0])
+      socket = SimpleH2Client.setup_connection(context)
+      port = context[:port]
+      SimpleH2Client.send_simple_headers(socket, 1, :get, "/body_response", port)
+      {:ok, 1, false, _, _} = SimpleH2Client.recv_headers(socket)
+      assert SimpleH2Client.recv_body(socket) == {:ok, 1, true, "OK"}
+    end
   end
 
   describe "settings exchange" do
@@ -3050,4 +3061,99 @@ defmodule HTTP2ProtocolTest do
       echo_components(conn)
     end
   end
+
+  describe "max_concurrent_streams flag" do
+    test "refuses streams when max_concurrent_streams limit is exceeded", context do
+      context =
+        context
+        |> https_server(http_2_options: [default_local_settings: [max_concurrent_streams: 2]])
+        |> Enum.into(context)
+
+      port = context.port
+
+      socket = setup_connection_with_custom_settings(context)
+
+      {:ok, send_ctx} =
+        SimpleH2Client.send_simple_headers(socket, 1, :get, "/slow_body_response", port)
+
+      {:ok, send_ctx} =
+        SimpleH2Client.send_simple_headers(socket, 3, :get, "/slow_body_response", port, send_ctx)
+
+      {:ok, send_ctx} =
+        SimpleH2Client.send_simple_headers(socket, 5, :get, "/slow_body_response", port, send_ctx)
+
+      t0 = for _ <- 1..3, do: SimpleH2Client.recv_frame(socket)
+
+      # Now fourth stream should be accepted since we're below the limit again
+      SimpleH2Client.send_simple_headers(socket, 7, :get, "/echo", context.port, send_ctx)
+
+      t1 = for _ <- 1..5, do: SimpleH2Client.recv_frame(socket)
+
+      t0
+      |> Enum.concat(t1)
+      |> assert
+      ~> in_any_order([
+        {:ok, :headers, integer(), 1, string()},
+        {:ok, :data, integer(), 1, "OK"},
+        {:ok, :headers, integer(), 3, string()},
+        {:ok, :data, integer(), 3, "OK"},
+        {:ok, :rst_stream, 0, 5, <<7::32>>},
+        {:ok, :headers, integer(), 7, string()},
+        {:ok, :data, integer(), 7, "OK"}
+      ])
+
+      assert Transport.recv(socket, 0) == {:error, :closed}
+    end
+
+    test "allows new streams after previous streams complete", context do
+      context =
+        context
+        |> https_server(http_2_options: [default_local_settings: [max_concurrent_streams: 2]])
+        |> Enum.into(context)
+
+      port = context.port
+
+      socket = setup_connection_with_custom_settings(context)
+
+      {:ok, send_ctx} =
+        SimpleH2Client.send_simple_headers(socket, 1, :get, "/slow_body_response", port)
+
+      {:ok, send_ctx} =
+        SimpleH2Client.send_simple_headers(socket, 3, :get, "/slow_body_response", port, send_ctx)
+
+      t0 = for _ <- 1..2, do: SimpleH2Client.recv_frame(socket)
+
+      SimpleH2Client.send_simple_headers(socket, 5, :get, "/slow_body_response", port, send_ctx)
+
+      t1 = for _ <- 1..5, do: SimpleH2Client.recv_frame(socket)
+
+      t0
+      |> Enum.concat(t1)
+      |> assert
+      ~> in_any_order([
+        {:ok, :headers, integer(), 1, string()},
+        {:ok, :headers, integer(), 3, string()},
+        {:ok, :headers, integer(), 5, string()},
+        {:ok, :data, integer(), 1, "OK"},
+        {:ok, :data, integer(), 3, "OK"},
+        {:ok, :data, integer(), 5, "OK"},
+        {:ok, :goaway, 0, 0, <<5::32, 0::32>>}
+      ])
+
+      assert Transport.recv(socket, 0) == {:error, :closed}
+    end
+
+    def slow_body_response(conn) do
+      Process.sleep(50)
+      conn |> send_resp(200, "OK")
+    end
+
+    # Helper function to set up connection when server has custom settings
+    defp setup_connection_with_custom_settings(context) do
+      socket = SimpleH2Client.tls_client(context)
+      SimpleH2Client.exchange_prefaces(socket, true)
+      SimpleH2Client.exchange_client_settings(socket)
+      socket
+    end
+  end
 end

test/support/simple_h2_client.ex

@@ -12,9 +12,16 @@ defmodule SimpleH2Client do
     socket
   end
 
-  def exchange_prefaces(socket) do
+  def exchange_prefaces(socket, with_settings \\ false) do
     Transport.send(socket, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
-    {:ok, <<0, 0, 0, 4, 0, 0, 0, 0, 0>>} = Transport.recv(socket, 9)
+    {:ok, <<length::24, 4, 0, 0, 0, 0, 0>>} = Transport.recv(socket, 9)
+
+    if with_settings and length > 0 do
+      {:ok, _settings_data} = Transport.recv(socket, length)
+    else
+      0 = length
+    end
+
     Transport.send(socket, <<0, 0, 0, 4, 1, 0, 0, 0, 0>>)
   end