From 746ad08705643a1e1966dcae2f343eecf25f2254 Mon Sep 17 00:00:00 2001
From: Mat Trudel <mat@geeky.net>
Date: Sat, 4 Jan 2025 13:08:29 -0500
Subject: [PATCH] Read a minimal number of bytes when sniffing for protocol
 (#449)

---
 lib/bandit/initial_handler.ex | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/lib/bandit/initial_handler.ex b/lib/bandit/initial_handler.ex
index 3bbff2f4..7e5ab273 100644
--- a/lib/bandit/initial_handler.ex
+++ b/lib/bandit/initial_handler.ex
@@ -57,19 +57,30 @@ defmodule Bandit.InitialHandler do
     end
   end
 
-  # Returns the protocol as suggested by received data, if possible
+  # Returns the protocol as suggested by received data, if possible.
+  # We do this in two phases so that we don't hang on *really* short HTTP/1
+  # requests that are less than 24 bytes
   @spec sniff_wire(ThousandIsland.Socket.t()) ::
           Bandit.HTTP2.Handler
           | :likely_tls
           | {:no_match, binary()}
           | {:error, :closed | :timeout | :inet.posix()}
   defp sniff_wire(socket) do
-    case ThousandIsland.Socket.recv(socket, 24) do
-      {:ok, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"} -> Bandit.HTTP2.Handler
-      {:ok, <<22::8, 3::8, minor::8, _::binary>>} when minor in [1, 3] -> :likely_tls
+    case ThousandIsland.Socket.recv(socket, 3) do
+      {:ok, "PRI" = buffer} -> sniff_wire_for_http2(socket, buffer)
+      {:ok, <<22::8, 3::8, minor::8>>} when minor in [1, 3] -> :likely_tls
       {:ok, data} -> {:no_match, data}
       {:error, :timeout} -> {:no_match, <<>>}
       {:error, error} -> {:error, error}
     end
   end
+
+  defp sniff_wire_for_http2(socket, buffer) do
+    case ThousandIsland.Socket.recv(socket, 21) do
+      {:ok, " * HTTP/2.0\r\n\r\nSM\r\n\r\n"} -> Bandit.HTTP2.Handler
+      {:ok, data} -> {:no_match, buffer <> data}
+      {:error, :timeout} -> {:no_match, buffer}
+      {:error, error} -> {:error, error}
+    end
+  end
 end