feat(http): rework Gun connection pooling

config/description.exs

@@ -3085,9 +3085,16 @@ config :pleroma, :config_description, [
       %{
         key: :max_connections,
         type: :integer,
-        description: "Maximum number of connections in the pool. Default: 250 connections.",
+        description:
+          "Maximum total number of connections in the pool. HTTP/1 origins may use multiple connections within this limit, while HTTP/2 connections are multiplexed. Default: 250 connections.",
         suggestions: [250]
       },
+      %{
+        key: :max_idle_time,
+        type: :integer,
+        description: "Time before an unused connection is closed. Default: 30000ms.",
+        suggestions: [30_000]
+      },
       %{
         key: :connect_timeout,
         type: :integer,

lib/pleroma/gun.ex

@@ -6,9 +6,12 @@ defmodule Pleroma.Gun do
   @callback open(charlist(), pos_integer(), map()) :: {:ok, pid()}
   @callback info(pid()) :: map()
   @callback close(pid()) :: :ok
+  @callback cancel(pid(), reference() | [reference()]) :: :ok
   @callback await_up(pid, pos_integer()) :: {:ok, atom()} | {:error, atom()}
   @callback connect(pid(), map()) :: reference()
-  @callback await(pid(), reference()) :: {:response, :fin, 200, []}
+  @callback await(pid(), reference()) :: tuple()
+  @callback await_tunnel_up(pid(), reference() | :undefined, timeout()) ::
+              {:ok, atom()} | {:error, term()}
   @callback set_owner(pid(), pid()) :: :ok
 
   defp api, do: Pleroma.Config.get([Pleroma.Gun], Pleroma.Gun.API)
@@ -19,11 +22,15 @@ defmodule Pleroma.Gun do
 
   def close(pid), do: api().close(pid)
 
+  def cancel(pid, stream), do: api().cancel(pid, stream)
+
   def await_up(pid, timeout \\ 5_000), do: api().await_up(pid, timeout)
 
   def connect(pid, opts), do: api().connect(pid, opts)
 
   def await(pid, ref), do: api().await(pid, ref)
 
+  def await_tunnel_up(pid, ref, timeout), do: api().await_tunnel_up(pid, ref, timeout)
+
   def set_owner(pid, owner), do: api().set_owner(pid, owner)
 end

lib/pleroma/gun/api.ex

@@ -32,6 +32,9 @@ defmodule Pleroma.Gun.API do
   @impl Gun
   defdelegate close(pid), to: :gun
 
+  @impl Gun
+  defdelegate cancel(pid, stream), to: :gun
+
   @impl Gun
   defdelegate await_up(pid, timeout \\ 5_000), to: :gun
 
@@ -41,6 +44,28 @@ defmodule Pleroma.Gun.API do
   @impl Gun
   defdelegate await(pid, ref), to: :gun
 
+  @impl Gun
+  def await_tunnel_up(pid, stream_ref, timeout) do
+    monitor_ref = Process.monitor(pid)
+
+    result =
+      receive do
+        {:gun_tunnel_up, ^pid, ^stream_ref, protocol} ->
+          {:ok, protocol}
+
+        {:gun_down, ^pid, _protocol, reason, _killed_streams} ->
+          {:error, {:down, reason}}
+
+        {:DOWN, ^monitor_ref, :process, ^pid, reason} ->
+          {:error, {:down, reason}}
+      after
+        timeout -> {:error, :timeout}
+      end
+
+    Process.demonitor(monitor_ref, [:flush])
+    result
+  end
+
   @impl Gun
   defdelegate set_owner(pid, owner), to: :gun
 end

lib/pleroma/gun/conn.ex

@@ -47,13 +47,14 @@ defmodule Pleroma.Gun.Conn do
       uri
       |> destination_opts()
       |> add_http2_opts(uri.scheme, Map.get(opts, :tls_opts, []))
+      |> add_proxy_auth(opts)
 
-    with open_opts <- Map.delete(opts, :tls_opts),
+    with open_opts <- proxy_open_opts(opts),
          {:ok, conn} <- Gun.open(proxy_host, proxy_port, open_opts),
-         {:ok, protocol} <- Gun.await_up(conn, opts[:connect_timeout]),
-         stream <- Gun.connect(conn, connect_opts),
-         {:response, :fin, 200, _} <- Gun.await(conn, stream) do
-      {:ok, conn, protocol}
+         {:ok, _proxy_protocol} <- await_up(conn, opts[:connect_timeout]) do
+      conn
+      |> Gun.connect(connect_opts)
+      |> await_proxy_connect(conn, opts[:connect_timeout])
     else
       error ->
         Logger.warning(
@@ -64,30 +65,29 @@ defmodule Pleroma.Gun.Conn do
     end
   end
 
-  defp do_open(uri, %{proxy: {proxy_type, proxy_host, proxy_port}} = opts) do
-    version =
-      proxy_type
-      |> to_string()
-      |> String.last()
-      |> case do
-        "4" -> 4
-        _ -> 5
+  defp do_open(_uri, %{proxy: {:socks4, _proxy_host, _proxy_port}}) do
+    {:error, :socks4_unsupported}
   end
 
+  defp do_open(uri, %{proxy: {proxy_type, proxy_host, proxy_port}} = opts)
+       when proxy_type in [:socks, :socks5] do
     socks_opts =
       uri
       |> destination_opts()
       |> add_http2_opts(uri.scheme, Map.get(opts, :tls_opts, []))
-      |> Map.put(:version, version)
+      |> Map.put(:version, 5)
+      |> add_socks_proxy_auth(opts)
 
-    opts =
+    open_opts =
       opts
+      |> proxy_open_opts()
       |> Map.put(:protocols, [:socks])
       |> Map.put(:socks_opts, socks_opts)
 
-    with {:ok, conn} <- Gun.open(proxy_host, proxy_port, opts),
-         {:ok, protocol} <- Gun.await_up(conn, opts[:connect_timeout]) do
-      {:ok, conn, protocol}
+    with {:ok, conn} <- Gun.open(proxy_host, proxy_port, open_opts),
+         {:ok, :socks} <- await_up(conn, opts[:connect_timeout]),
+         {:ok, protocol} <- await_tunnel_up(conn, :undefined, opts[:connect_timeout]) do
+      {:ok, conn, protocol, nil}
     else
       error ->
         Logger.warning(
@@ -100,10 +100,11 @@ defmodule Pleroma.Gun.Conn do
 
   defp do_open(%URI{host: host, port: port} = uri, opts) do
     host = Pleroma.HTTP.AdapterHelper.parse_host(host)
+    opts = Map.put(opts, :transport, transport(uri.scheme))
 
     with {:ok, conn} <- Gun.open(host, port, opts),
-         {:ok, protocol} <- Gun.await_up(conn, opts[:connect_timeout]) do
-      {:ok, conn, protocol}
+         {:ok, protocol} <- await_up(conn, opts[:connect_timeout]) do
+      {:ok, conn, protocol, nil}
     else
       error ->
         Logger.warning(
@@ -120,11 +121,90 @@ defmodule Pleroma.Gun.Conn do
   end
 
   defp add_http2_opts(opts, "https", tls_opts) do
-    Map.merge(opts, %{protocols: [:http2], transport: :tls, tls_opts: tls_opts})
+    Map.merge(opts, %{protocols: [:http2, :http], transport: :tls, tls_opts: tls_opts})
   end
 
   defp add_http2_opts(opts, _, _), do: opts
 
+  defp add_proxy_auth(connect_opts, %{proxy_auth: {username, password}})
+       when is_binary(username) and is_binary(password) do
+    Map.merge(connect_opts, %{username: username, password: password})
+  end
+
+  defp add_proxy_auth(connect_opts, _), do: connect_opts
+
+  defp add_socks_proxy_auth(socks_opts, %{proxy_auth: {username, password}})
+       when is_binary(username) and is_binary(password) do
+    Map.put(socks_opts, :auth, [{:username_password, username, password}])
+  end
+
+  defp add_socks_proxy_auth(socks_opts, _), do: socks_opts
+
+  defp proxy_open_opts(opts) do
+    proxy_tls_opts = Map.get(opts, :proxy_tls_opts, [])
+
+    opts =
+      opts
+      |> Map.drop([:proxy, :proxy_auth, :proxy_tls_opts, :tls_opts])
+      |> Map.put_new(:transport, :tcp)
+
+    if opts.transport == :tls do
+      Map.put(opts, :tls_opts, proxy_tls_opts)
+    else
+      opts
+    end
+  end
+
+  defp await_up(conn, timeout) do
+    case Gun.await_up(conn, timeout) do
+      {:ok, protocol} ->
+        {:ok, protocol}
+
+      error ->
+        Gun.close(conn)
+        error
+    end
+  end
+
+  defp await_proxy_connect(stream, conn, timeout) do
+    case Gun.await(conn, stream) do
+      {:response, :fin, 200, _headers} ->
+        case await_tunnel_up(conn, stream, timeout) do
+          {:ok, protocol} -> {:ok, conn, protocol, stream}
+          error -> error
+        end
+
+      {:response, _fin, 403, _headers} ->
+        close_with_error(conn, :unauthorized)
+
+      {:response, _fin, 407, _headers} ->
+        close_with_error(conn, :proxy_auth_failed)
+
+      error ->
+        Gun.close(conn)
+        error
+    end
+  end
+
+  defp close_with_error(conn, reason) do
+    Gun.close(conn)
+    {:error, reason}
+  end
+
+  defp transport("https"), do: :tls
+  defp transport(_), do: :tcp
+
+  defp await_tunnel_up(conn, stream_ref, timeout) do
+    case Gun.await_tunnel_up(conn, stream_ref, timeout) do
+      {:ok, protocol} ->
+        {:ok, protocol}
+
+      error ->
+        Gun.close(conn)
+        error
+    end
+  end
+
   def compose_uri_log(%URI{scheme: scheme, host: host, path: path}) do
     "#{scheme}://#{host}#{path}"
   end

lib/pleroma/gun/connection_pool.ex

@@ -4,53 +4,132 @@
 
 defmodule Pleroma.Gun.ConnectionPool do
   @registry __MODULE__
+  @locks Module.concat(__MODULE__, Locks)
 
   alias Pleroma.Gun.ConnectionPool.WorkerSupervisor
 
+  @connection_option_keys [
+    :http_opts,
+    :http2_opts,
+    :protocols,
+    :proxy,
+    :proxy_auth,
+    :proxy_tls_opts,
+    :retry,
+    :retry_timeout,
+    :socks_opts,
+    :tcp_opts,
+    :tls_opts,
+    :transport,
+    :ws_opts
+  ]
+
   def children do
     [
-      {Registry, keys: :unique, name: @registry},
+      {Registry, keys: :duplicate, name: @registry},
+      {Registry, keys: :unique, name: @locks},
       Pleroma.Gun.ConnectionPool.WorkerSupervisor
     ]
   end
 
   @spec get_conn(URI.t(), keyword()) :: {:ok, pid()} | {:error, term()}
   def get_conn(uri, opts) do
-    key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
+    key = connection_key(uri, opts)
 
-    case Registry.lookup(@registry, key) do
-      # The key has already been registered, but connection is not up yet
-      [{worker_pid, nil}] ->
-        get_gun_pid_from_worker(worker_pid, true)
-
-      [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
-        GenServer.call(worker_pid, :add_client)
+    case checkout_existing(Registry.lookup(@registry, key)) do
+      {:ok, gun_pid} ->
         {:ok, gun_pid}
 
-      [] ->
-        # :gun.set_owner fails in :connected state for whatevever reason,
-        # so we open the connection in the process directly and send it's pid back
-        # We trust gun to handle timeouts by itself
+      :not_available ->
+        provision_or_wait(key, uri, opts)
+    end
+  end
+
+  defp checkout_existing([]), do: :not_available
+
+  defp checkout_existing([{_worker_pid, nil} | workers]), do: checkout_existing(workers)
+
+  defp checkout_existing([{worker_pid, _value} | workers]) do
+    case checkout_worker(worker_pid) do
+      {:ok, gun_pid} -> {:ok, gun_pid}
+      _ -> checkout_existing(workers)
+    end
+  end
+
+  defp provision_or_wait(key, uri, opts) do
+    result = with_route_lock(key, fn -> provision(key, uri, opts) end)
+
+    case result do
+      {:ok, gun_pid} ->
+        {:ok, gun_pid}
+
+      {:new, worker_pid} ->
+        get_gun_pid_from_worker(worker_pid)
+
+      {:wait, worker_pid} ->
+        case checkout_worker(worker_pid, :infinity) do
+          {:ok, gun_pid} -> {:ok, gun_pid}
+          :busy -> get_conn(uri, opts)
+          error -> error
+        end
+
+      error ->
+        error
+    end
+  end
+
+  defp provision(key, uri, opts) do
+    workers = Registry.lookup(@registry, key)
+
+    case checkout_existing(workers) do
+      {:ok, gun_pid} ->
+        {:ok, gun_pid}
+
+      :not_available ->
+        case Enum.find(workers, fn {_worker_pid, value} -> is_nil(value) end) do
+          {worker_pid, nil} -> {:wait, worker_pid}
+          nil -> start_worker(key, uri, opts)
+        end
+    end
+  end
+
+  defp with_route_lock(key, function) do
+    case Registry.register(@locks, key, nil) do
+      {:ok, _} ->
+        try do
+          function.()
+        after
+          Registry.unregister(@locks, key)
+        end
+
+      {:error, {:already_registered, _owner}} ->
+        Process.sleep(1)
+        with_route_lock(key, function)
+    end
+  end
+
+  defp start_worker(key, uri, opts) do
+    # :gun.set_owner fails in :connected state, so the worker opens the connection itself.
     case WorkerSupervisor.start_worker([key, uri, opts, self()]) do
-          {:ok, worker_pid} ->
-            get_gun_pid_from_worker(worker_pid, false)
-
-          {:error, {:already_started, worker_pid}} ->
-            get_gun_pid_from_worker(worker_pid, true)
-
-          err ->
-            err
-        end
+      {:ok, worker_pid} -> {:new, worker_pid}
+      error -> error
     end
   end
 
-  defp get_gun_pid_from_worker(worker_pid, register) do
+  defp checkout_worker(worker_pid, timeout \\ 5_000) do
+    try do
+      GenServer.call(worker_pid, :checkout, timeout)
+    catch
+      :exit, reason -> {:error, {:connection_worker_exit, reason}}
+    end
+  end
+
+  defp get_gun_pid_from_worker(worker_pid) do
     # GenServer.call will block the process for timeout length if
     # the server crashes on startup (which will happen if gun fails to connect)
     # so instead we use cast + monitor
 
     ref = Process.monitor(worker_pid)
-    if register, do: GenServer.cast(worker_pid, {:add_client, self()})
 
     receive do
       {:conn_pid, pid} ->
@@ -68,19 +147,115 @@ defmodule Pleroma.Gun.ConnectionPool do
 
   @spec release_conn(pid()) :: :ok
   def release_conn(conn_pid) do
-    # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _, _, _}}} when gun_pid == conn_pid ->
+    # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _tunnel}}} when gun_pid == conn_pid ->
     #    worker_pid end)
     query_result =
       Registry.select(@registry, [
-        {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
+        {{:_, :"$1", {:"$2", :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
       ])
 
     case query_result do
       [worker_pid] ->
+        try do
           GenServer.call(worker_pid, :remove_client)
+        catch
+          :exit, _ -> :ok
+        end
 
       [] ->
         :ok
     end
   end
+
+  @spec discard_conn(pid()) :: :ok
+  def discard_conn(conn_pid) do
+    case worker_for_conn(conn_pid) do
+      nil ->
+        :ok
+
+      worker_pid ->
+        DynamicSupervisor.terminate_child(WorkerSupervisor, worker_pid)
+    end
+  end
+
+  @spec cancel_stream(pid(), reference() | [reference()]) :: :ok
+  def cancel_stream(conn_pid, stream) do
+    case worker_for_conn(conn_pid) do
+      nil ->
+        Pleroma.Gun.cancel(conn_pid, stream)
+
+      worker_pid ->
+        try do
+          GenServer.call(worker_pid, {:cancel_stream, stream})
+        catch
+          :exit, _ -> :ok
+        end
+    end
+  end
+
+  @spec register_stream(pid(), reference() | [reference()]) :: :ok
+  def register_stream(conn_pid, stream) do
+    call_worker(conn_pid, {:register_stream, stream})
+  end
+
+  @spec finish_stream(pid(), reference() | [reference()]) :: :ok
+  def finish_stream(conn_pid, stream) do
+    call_worker(conn_pid, {:finish_stream, stream})
+  end
+
+  @spec release_stream(pid(), reference() | [reference()]) :: :ok
+  def release_stream(conn_pid, stream) do
+    call_worker(conn_pid, {:release_stream, stream})
+  end
+
+  @spec tunnel_ref(pid()) :: reference() | nil
+  def tunnel_ref(conn_pid) do
+    case Registry.select(@registry, [
+           {{:_, :_, {:"$1", :"$2"}}, [{:==, :"$1", conn_pid}], [:"$2"]}
+         ]) do
+      [tunnel] -> tunnel
+      [] -> nil
+    end
+  end
+
+  defp worker_for_conn(conn_pid) do
+    case Registry.select(@registry, [
+           {{:_, :"$1", {:"$2", :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
+         ]) do
+      [worker_pid] -> worker_pid
+      [] -> nil
+    end
+  end
+
+  defp call_worker(conn_pid, message) do
+    case worker_for_conn(conn_pid) do
+      nil ->
+        :ok
+
+      worker_pid ->
+        try do
+          GenServer.call(worker_pid, message)
+        catch
+          :exit, _ -> :ok
+        end
+    end
+  end
+
+  defp connection_key(uri, opts) do
+    origin = {
+      String.downcase(uri.scheme),
+      String.downcase(uri.host),
+      uri.port
+    }
+
+    profile =
+      opts
+      |> Map.new()
+      |> Map.take(@connection_option_keys)
+      |> :erlang.term_to_binary()
+      |> then(&:crypto.hash(:sha256, &1))
+
+    fingerprint = Base.encode16(profile, case: :lower)
+    Enum.join(Tuple.to_list(origin) ++ [fingerprint], ":")
+  end
 end

lib/pleroma/gun/connection_pool/reclaimer.ex

@@ -9,11 +9,11 @@ defmodule Pleroma.Gun.ConnectionPool.Reclaimer do
 
   def start_monitor do
     pid =
-      case GenServer.start(__MODULE__, [], name: {:via, Registry, {registry(), "reclaimer"}}) do
+      case GenServer.start(__MODULE__, [], name: __MODULE__) do
         {:ok, pid} ->
           pid
 
-        {:error, {:already_registered, pid}} ->
+        {:error, {:already_started, pid}} ->
           pid
       end
 
@@ -41,17 +41,16 @@ defmodule Pleroma.Gun.ConnectionPool.Reclaimer do
       reclaim_max: reclaim_max
     })
 
-    # :ets.fun2ms(
-    # fn {_, {worker_pid, {_, used_by, crf, last_reference}}} when used_by == [] ->
-    #   {worker_pid, crf, last_reference} end)
-    unused_conns =
+    workers =
       Registry.select(
         registry(),
         [
-          {{:_, :"$1", {:_, :"$2", :"$3", :"$4"}}, [{:==, :"$2", []}], [{{:"$1", :"$3", :"$4"}}]}
+          {{:_, :"$1", {:_, :_}}, [], [:"$1"]}
         ]
       )
 
+    unused_conns = Enum.flat_map(workers, &reclaim_info/1)
+
     case unused_conns do
       [] ->
         :telemetry.execute(
@@ -65,25 +64,43 @@ defmodule Pleroma.Gun.ConnectionPool.Reclaimer do
         {:stop, :no_unused_conns, nil}
 
       unused_conns ->
-        reclaimed =
+        reclaimed_count =
           unused_conns
-          |> Enum.sort(fn {_pid1, crf1, last_reference1}, {_pid2, crf2, last_reference2} ->
-            crf1 <= crf2 and last_reference1 <= last_reference2
-          end)
-          |> Enum.take(reclaim_max)
+          |> Enum.sort_by(fn {_pid, crf, last_reference} -> {crf, last_reference} end)
+          |> Enum.reduce_while(0, fn
+            _worker, count when count == reclaim_max ->
+              {:halt, count}
 
-        reclaimed
-        |> Enum.each(fn {pid, _, _} ->
-          DynamicSupervisor.terminate_child(Pleroma.Gun.ConnectionPool.WorkerSupervisor, pid)
+            {worker_pid, _crf, _last_reference}, count ->
+              if reclaim(worker_pid), do: {:cont, count + 1}, else: {:cont, count}
           end)
 
         :telemetry.execute(
           [:pleroma, :connection_pool, :reclaim, :stop],
-          %{reclaimed_count: Enum.count(reclaimed)},
+          %{reclaimed_count: reclaimed_count},
           %{max_connections: max_connections}
         )
 
         {:stop, :normal, nil}
     end
   end
+
+  defp reclaim_info(worker_pid) do
+    try do
+      case GenServer.call(worker_pid, :reclaim_info) do
+        {:idle, crf, last_reference} -> [{worker_pid, crf, last_reference}]
+        :in_use -> []
+      end
+    catch
+      :exit, _ -> []
+    end
+  end
+
+  defp reclaim(worker_pid) do
+    try do
+      GenServer.call(worker_pid, :reclaim) == :reclaimed
+    catch
+      :exit, _ -> false
+    end
+  end
 end

lib/pleroma/gun/connection_pool/worker.ex

@@ -8,25 +8,22 @@ defmodule Pleroma.Gun.ConnectionPool.Worker do
 
   defp registry, do: Pleroma.Gun.ConnectionPool
 
-  def start_link([key | _] = opts) do
-    GenServer.start_link(__MODULE__, opts, name: {:via, Registry, {registry(), key}})
-  end
+  def start_link(opts), do: GenServer.start_link(__MODULE__, opts)
 
   @impl true
-  def init([_key, _uri, _opts, _client_pid] = opts) do
+  def init([key, _uri, _opts, _client_pid] = opts) do
+    {:ok, _} = Registry.register(registry(), key, nil)
     {:ok, nil, {:continue, {:connect, opts}}}
   end
 
   @impl true
   def handle_continue({:connect, [key, uri, opts, client_pid]}, _) do
-    with {:ok, conn_pid, protocol} <- Gun.Conn.open(uri, opts),
+    with {:ok, conn_pid, protocol, tunnel} <- Gun.Conn.open(uri, opts),
          Process.link(conn_pid) do
       time = :erlang.monotonic_time(:millisecond)
 
-      {_, _} =
-        Registry.update_value(registry(), key, fn _ ->
-          {conn_pid, [client_pid], 1, time}
-        end)
+      Registry.unregister(registry(), key)
+      {:ok, _} = Registry.register(registry(), key, {conn_pid, tunnel})
 
       send(client_pid, {:conn_pid, conn_pid})
 
@@ -34,8 +31,14 @@ defmodule Pleroma.Gun.ConnectionPool.Worker do
        %{
          key: key,
          timer: nil,
-         client_monitors: %{client_pid => Process.monitor(client_pid)},
-         protocol: protocol
+         clients: %{
+           client_pid => %{count: 1, monitor: Process.monitor(client_pid), streams: MapSet.new()}
+         },
+         conn_pid: conn_pid,
+         protocol: protocol,
+         proxied: Keyword.has_key?(opts, :proxy),
+         crf: 1,
+         last_reference: time
        }, :hibernate}
     else
       err ->
@@ -44,106 +47,225 @@ defmodule Pleroma.Gun.ConnectionPool.Worker do
   end
 
   @impl true
-  def handle_cast({:add_client, client_pid}, state) do
-    case handle_call(:add_client, {client_pid, nil}, state) do
-      {:reply, conn_pid, state, :hibernate} ->
-        send(client_pid, {:conn_pid, conn_pid})
-        {:noreply, state, :hibernate}
-    end
+  def handle_call(:checkout, _from, %{protocol: protocol, clients: clients} = state)
+      when protocol != :http2 and map_size(clients) > 0 do
+    {:reply, :busy, state, :hibernate}
   end
 
-  @impl true
-  def handle_cast({:remove_client, client_pid}, state) do
-    case handle_call(:remove_client, {client_pid, nil}, state) do
-      {:reply, _, state, :hibernate} ->
-        {:noreply, state, :hibernate}
-    end
-  end
-
-  @impl true
-  def handle_call(:add_client, {client_pid, _}, %{key: key, protocol: protocol} = state) do
+  def handle_call(:checkout, {client_pid, _}, %{key: key, protocol: protocol} = state) do
     time = :erlang.monotonic_time(:millisecond)
 
-    {{conn_pid, used_by, _, _}, _} =
-      Registry.update_value(registry(), key, fn {conn_pid, used_by, crf, last_reference} ->
-        {conn_pid, [client_pid | used_by], crf(time - last_reference, crf), time}
-      end)
+    [{_, {conn_pid, _tunnel}}] =
+      Registry.lookup(registry(), key) |> Enum.filter(&(elem(&1, 0) == self()))
+
+    state =
+      state
+      |> cancel_idle_timer()
+      |> add_client(client_pid)
+      |> Map.put(:crf, crf(time - state.last_reference, state.crf))
+      |> Map.put(:last_reference, time)
 
     :telemetry.execute(
       [:pleroma, :connection_pool, :client, :add],
-      %{client_pid: client_pid, clients: used_by},
+      %{client_pid: client_pid, clients: Map.keys(state.clients)},
       %{key: state.key, protocol: protocol}
     )
 
-    state =
-      if state.timer != nil do
-        Process.cancel_timer(state[:timer])
-        %{state | timer: nil}
-      else
-        state
-      end
-
-    ref = Process.monitor(client_pid)
-
-    state = put_in(state.client_monitors[client_pid], ref)
-    {:reply, conn_pid, state, :hibernate}
+    {:reply, {:ok, conn_pid}, state, :hibernate}
   end
 
   @impl true
-  def handle_call(:remove_client, {client_pid, _}, %{key: key} = state) do
-    {{_conn_pid, used_by, _crf, _last_reference}, _} =
-      Registry.update_value(registry(), key, fn {conn_pid, used_by, crf, last_reference} ->
-        {conn_pid, List.delete(used_by, client_pid), crf, last_reference}
-      end)
+  def handle_call(:remove_client, {client_pid, _}, state) do
+    case active_streams(state, client_pid) do
+      [] ->
+        {:reply, :ok, remove_client(state, client_pid, true), :hibernate}
 
-    {ref, state} = pop_in(state.client_monitors[client_pid])
+      streams when state.protocol == :http2 ->
+        Enum.each(streams, &(:ok = Gun.cancel(state.conn_pid, &1)))
+        {:reply, :ok, remove_client(state, client_pid, true), :hibernate}
 
-    Process.demonitor(ref, [:flush])
-
-    timer =
-      if used_by == [] do
-        max_idle = Pleroma.Config.get([:connections_pool, :max_idle_time], 30_000)
-        Process.send_after(self(), :idle_close, max_idle)
-      else
-        nil
+      _streams ->
+        {:stop, :normal, :ok, state}
+    end
   end
 
-    {:reply, :ok, %{state | timer: timer}, :hibernate}
+  def handle_call({:register_stream, stream}, {client_pid, _}, state) do
+    {:reply, :ok, update_stream(state, client_pid, &MapSet.put(&1, stream)), :hibernate}
+  end
+
+  def handle_call({:finish_stream, stream}, {client_pid, _}, state) do
+    {:reply, :ok, update_stream(state, client_pid, &MapSet.delete(&1, stream)), :hibernate}
+  end
+
+  def handle_call({:release_stream, stream}, {client_pid, _}, state) do
+    state = update_stream(state, client_pid, &MapSet.delete(&1, stream))
+    {:reply, :ok, remove_client(state, client_pid, true), :hibernate}
+  end
+
+  def handle_call({:cancel_stream, stream}, {client_pid, _}, %{protocol: :http2} = state) do
+    :ok = Gun.cancel(state.conn_pid, stream)
+    state = update_stream(state, client_pid, &MapSet.delete(&1, stream))
+    {:reply, :ok, remove_client(state, client_pid, true), :hibernate}
+  end
+
+  def handle_call({:cancel_stream, _stream}, _from, state) do
+    {:stop, :normal, :ok, state}
+  end
+
+  def handle_call(:reclaim_info, _from, %{clients: clients} = state)
+      when map_size(clients) == 0 do
+    {:reply, {:idle, state.crf, state.last_reference}, state, :hibernate}
+  end
+
+  def handle_call(:reclaim_info, _from, state) do
+    {:reply, :in_use, state, :hibernate}
+  end
+
+  def handle_call(:reclaim, _from, %{clients: clients} = state) when map_size(clients) == 0 do
+    {:stop, :normal, :reclaimed, state}
+  end
+
+  def handle_call(:reclaim, _from, state) do
+    {:reply, :in_use, state, :hibernate}
   end
 
   @impl true
-  def handle_info(:idle_close, state) do
-    # Gun monitors the owner process, and will close the connection automatically
-    # when it's terminated
+  def handle_info({:idle_close, token}, %{timer: {_timer_ref, token}, clients: clients} = state)
+      when map_size(clients) == 0 do
+    # Gun monitors the owner process, and will close the connection automatically.
     {:stop, :normal, state}
   end
 
+  def handle_info({:idle_close, _token}, state), do: {:noreply, state, :hibernate}
+
   @impl true
-  def handle_info({:gun_up, _pid, _protocol}, state) do
-    {:noreply, state, :hibernate}
+  def handle_info({:gun_up, _pid, protocol}, state) do
+    {:noreply, %{state | protocol: protocol}, :hibernate}
   end
 
-  # Gracefully shutdown if the connection got closed without any streams left
+  @impl true
+  def handle_info({:gun_error, conn_pid, stream, _reason}, %{conn_pid: conn_pid} = state) do
+    {:noreply, delete_stream(state, stream), :hibernate}
+  end
+
+  @impl true
+  def handle_info({:gun_down, _pid, _protocol, _reason, _streams}, %{proxied: true} = state) do
+    {:stop, :normal, state}
+  end
+
+  # Gracefully shutdown if the connection got closed without any streams left.
   @impl true
   def handle_info({:gun_down, _pid, _protocol, _reason, []}, state) do
     {:stop, :normal, state}
   end
 
-  # Otherwise, wait for retry
+  # Otherwise, wait for retry.
   @impl true
   def handle_info({:gun_down, _pid, _protocol, _reason, _killed_streams}, state) do
     {:noreply, state, :hibernate}
   end
 
   @impl true
-  def handle_info({:DOWN, _ref, :process, pid, reason}, state) do
+  def handle_info({:DOWN, ref, :process, pid, reason}, state) do
+    case state.clients[pid] do
+      %{monitor: ^ref} ->
         :telemetry.execute(
           [:pleroma, :connection_pool, :client, :dead],
           %{client_pid: pid, reason: reason},
           %{key: state.key}
         )
 
-    handle_cast({:remove_client, pid}, state)
+        handle_dead_client(state, pid)
+
+      _ ->
+        {:noreply, state, :hibernate}
+    end
+  end
+
+  defp add_client(state, client_pid) do
+    update_in(state.clients, fn clients ->
+      case clients[client_pid] do
+        nil ->
+          Map.put(clients, client_pid, %{
+            count: 1,
+            monitor: Process.monitor(client_pid),
+            streams: MapSet.new()
+          })
+
+        client ->
+          Map.put(clients, client_pid, %{client | count: client.count + 1})
+      end
+    end)
+  end
+
+  defp handle_dead_client(state, client_pid) do
+    streams = active_streams(state, client_pid)
+
+    cond do
+      streams == [] ->
+        {:noreply, remove_client(state, client_pid, false, true), :hibernate}
+
+      state.protocol == :http2 ->
+        Enum.each(streams, &(:ok = Gun.cancel(state.conn_pid, &1)))
+        {:noreply, remove_client(state, client_pid, false, true), :hibernate}
+
+      true ->
+        {:stop, :normal, state}
+    end
+  end
+
+  defp active_streams(state, client_pid) do
+    case state.clients[client_pid] do
+      %{streams: streams} -> MapSet.to_list(streams)
+      nil -> []
+    end
+  end
+
+  defp update_stream(state, client_pid, update) do
+    update_in(state.clients, fn clients ->
+      case clients[client_pid] do
+        nil -> clients
+        client -> Map.put(clients, client_pid, %{client | streams: update.(client.streams)})
+      end
+    end)
+  end
+
+  defp delete_stream(state, stream) do
+    update_in(state.clients, fn clients ->
+      Map.new(clients, fn {pid, client} ->
+        {pid, %{client | streams: MapSet.delete(client.streams, stream)}}
+      end)
+    end)
+  end
+
+  defp remove_client(state, client_pid, demonitor, all \\ false) do
+    case state.clients[client_pid] do
+      nil ->
+        state
+
+      %{count: count} = client when count > 1 and not all ->
+        put_in(state.clients[client_pid], %{client | count: count - 1})
+
+      client ->
+        if demonitor, do: Process.demonitor(client.monitor, [:flush])
+        state = update_in(state.clients, &Map.delete(&1, client_pid))
+
+        if map_size(state.clients) == 0, do: schedule_idle_close(state), else: state
+    end
+  end
+
+  defp cancel_idle_timer(%{timer: nil} = state), do: state
+
+  defp cancel_idle_timer(%{timer: {timer_ref, _token}} = state) do
+    Process.cancel_timer(timer_ref)
+    %{state | timer: nil}
+  end
+
+  defp schedule_idle_close(state) do
+    max_idle = Pleroma.Config.get([:connections_pool, :max_idle_time], 30_000)
+    token = make_ref()
+    timer_ref = Process.send_after(self(), {:idle_close, token}, max_idle)
+    %{state | timer: {timer_ref, token}}
   end
 
   # LRFU policy: https://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.55.1478

lib/pleroma/gun/connection_pool/worker_supervisor.ex

@@ -26,7 +26,14 @@ defmodule Pleroma.Gun.ConnectionPool.WorkerSupervisor do
   def start_worker(opts, true) do
     case DynamicSupervisor.start_child(__MODULE__, {Worker, opts}) do
       {:error, :max_children} ->
-        :telemetry.execute([:pleroma, :connection_pool, :provision_failure], %{opts: opts})
+        [key | _] = opts
+
+        :telemetry.execute(
+          [:pleroma, :connection_pool, :provision_failure],
+          %{},
+          %{key: key}
+        )
+
         {:error, :pool_full}
 
       res ->

lib/pleroma/http/adapter/gun.ex

@@ -0,0 +1,207 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Pleroma.HTTP.Adapter.Gun do
+  @behaviour Tesla.Adapter
+
+  alias Pleroma.Gun.ConnectionPool
+  alias Tesla.Multipart
+
+  @default_timeout 1_000
+
+  @impl Tesla.Adapter
+  def call(env, opts) do
+    adapter_opts =
+      Tesla.Adapter.opts(
+        [close_conn: true, body_as: :plain, send_body: :at_once, receive: true],
+        env,
+        opts
+      )
+      |> Map.new()
+
+    with {:ok, status, headers, body} <- request(env, adapter_opts) do
+      env = %{env | status: status, headers: format_headers(headers), body: body}
+      {:ok, put_stream_ref(env, body)}
+    end
+  end
+
+  defp put_stream_ref(env, %{stream: stream}) do
+    put_in(env.opts[:adapter][:stream], stream)
+  end
+
+  defp put_stream_ref(env, _body), do: env
+
+  defp request(env, %{conn: conn, tunnel: tunnel} = opts) do
+    uri = URI.parse(Tesla.build_url(env))
+    path = Tesla.Adapter.Shared.prepare_path(uri.path, uri.query)
+    method = Tesla.Adapter.Shared.format_method(env.method)
+    {headers, body, send_body} = prepare_body(env.headers, env.body, opts[:send_body])
+    request_opts = request_opts(opts[:reply_to], tunnel)
+    stream = open_stream(conn, method, path, headers, body, request_opts, send_body)
+    :ok = ConnectionPool.register_stream(conn, stream)
+    :ok = send_stream_body(conn, stream, body, send_body)
+    response = read_response(conn, stream, opts)
+
+    case response do
+      {:error, reason} ->
+        ConnectionPool.cancel_stream(conn, stream)
+        {:error, {:gun_stream_error, reason}}
+
+      response ->
+        unless opts[:body_as] in [:chunks, :stream] and stream_body?(response) do
+          :ok = ConnectionPool.finish_stream(conn, stream)
+        end
+
+        if opts[:close_conn] and opts[:body_as] not in [:stream, :chunks] do
+          :gun.close(conn)
+        end
+
+        response
+    end
+  end
+
+  defp stream_body?({:ok, _status, _headers, %{stream: _stream}}), do: true
+  defp stream_body?({:ok, _status, _headers, %Stream{}}), do: true
+  defp stream_body?(_response), do: false
+
+  defp request_opts(reply_to, nil), do: %{reply_to: reply_to || self()}
+
+  defp request_opts(reply_to, tunnel) do
+    %{reply_to: reply_to || self(), tunnel: tunnel}
+  end
+
+  defp prepare_body(headers, %Multipart{} = multipart, _send_body) do
+    {format_headers(headers ++ Multipart.headers(multipart)), Multipart.body(multipart), :stream}
+  end
+
+  defp prepare_body(headers, body, _send_body)
+       when is_struct(body, Stream) or is_function(body) do
+    {format_headers(headers), body, :stream}
+  end
+
+  defp prepare_body(headers, body, send_body) do
+    {format_headers(headers), body || "", send_body}
+  end
+
+  defp open_stream(conn, method, path, headers, _body, request_opts, :stream) do
+    :gun.headers(conn, method, path, headers, request_opts)
+  end
+
+  defp open_stream(conn, method, path, headers, body, request_opts, :at_once) do
+    :gun.request(conn, method, path, headers, body, request_opts)
+  end
+
+  defp send_stream_body(conn, stream, body, :stream) do
+    try do
+      for data <- body, do: :ok = :gun.data(conn, stream, :nofin, data)
+      :gun.data(conn, stream, :fin, "")
+    catch
+      kind, reason ->
+        ConnectionPool.cancel_stream(conn, stream)
+        :erlang.raise(kind, reason, __STACKTRACE__)
+    end
+  end
+
+  defp send_stream_body(_conn, _stream, _body, :at_once), do: :ok
+
+  defp read_response(conn, stream, opts) do
+    receive? = opts[:receive]
+
+    receive do
+      {:gun_response, ^conn, ^stream, :fin, status, headers} ->
+        {:ok, status, headers, ""}
+
+      {:gun_response, ^conn, ^stream, :nofin, status, headers} ->
+        format_response(conn, stream, opts, status, headers, opts[:body_as])
+
+      {:gun_up, ^conn, _protocol} when receive? ->
+        read_response(conn, stream, opts)
+
+      {:gun_error, ^conn, reason} ->
+        {:error, reason}
+
+      {:gun_error, ^conn, ^stream, reason} ->
+        {:error, reason}
+
+      {:gun_down, ^conn, _protocol, _reason, _killed_streams} when receive? ->
+        read_response(conn, stream, opts)
+
+      {:DOWN, _ref, :process, ^conn, reason} ->
+        {:error, reason}
+    after
+      opts[:timeout] || @default_timeout -> {:error, :recv_response_timeout}
+    end
+  end
+
+  defp format_response(conn, stream, opts, status, headers, :plain) do
+    case read_body(conn, stream, opts) do
+      {:ok, body} ->
+        {:ok, status, headers, body}
+
+      {:error, error} ->
+        :ok = :gun.flush(stream)
+        {:error, error}
+    end
+  end
+
+  defp format_response(conn, stream, opts, status, headers, :stream) do
+    body =
+      Stream.resource(
+        fn -> :reading end,
+        fn
+          :reading ->
+            case Tesla.Adapter.Gun.read_chunk(conn, stream, opts) do
+              {:nofin, part} -> {[part], :reading}
+              {:fin, part} -> {[part], :done}
+              {:error, reason} -> raise "Gun stream failed: #{inspect(reason)}"
+            end
+
+          :done ->
+            {:halt, :done}
+        end,
+        fn _state ->
+          if opts[:close_conn], do: :gun.close(conn)
+        end
+      )
+
+    {:ok, status, headers, body}
+  end
+
+  defp format_response(conn, stream, opts, status, headers, :chunks) do
+    {:ok, status, headers, %{pid: conn, stream: stream, opts: Map.to_list(opts)}}
+  end
+
+  defp read_body(conn, stream, opts, acc \\ "") do
+    receive do
+      {:gun_data, ^conn, ^stream, :fin, body} ->
+        append_body(acc, body, opts[:max_body])
+
+      {:gun_data, ^conn, ^stream, :nofin, body} ->
+        with {:ok, acc} <- append_body(acc, body, opts[:max_body]) do
+          read_body(conn, stream, opts, acc)
+        end
+
+      {:gun_error, ^conn, ^stream, reason} ->
+        {:error, reason}
+
+      {:DOWN, _ref, :process, ^conn, reason} ->
+        {:error, reason}
+    after
+      opts[:timeout] || @default_timeout -> {:error, :recv_body_timeout}
+    end
+  end
+
+  defp append_body(acc, part, nil), do: {:ok, acc <> part}
+
+  defp append_body(acc, part, limit) do
+    body = acc <> part
+    if byte_size(body) <= limit, do: {:ok, body}, else: {:error, :body_too_large}
+  end
+
+  defp format_headers(headers) do
+    Enum.map(headers, fn {key, value} ->
+      {String.downcase(to_string(key)), to_string(value)}
+    end)
+  end
+end

lib/pleroma/http/adapter_helper/gun.ex

@@ -48,10 +48,10 @@ defmodule Pleroma.HTTP.AdapterHelper.Gun do
     Keyword.put(opts, :timeout, recv_timeout)
   end
 
-  # Gun uses [body_as: :stream]
+  # Keep the pool lease until the caller consumes or cancels the body.
   defp maybe_stream(opts) do
     case Keyword.pop(opts, :stream, nil) do
-      {true, opts} -> Keyword.put(opts, :body_as, :stream)
+      {true, opts} -> Keyword.put(opts, :body_as, :chunks)
       {_, opts} -> opts
     end
   end

lib/pleroma/reverse_proxy.ex

@@ -110,12 +110,8 @@ defmodule Pleroma.ReverseProxy do
       end
 
     with {:ok, nil} <- @cachex.get(:failed_proxy_url_cache, url),
-         {:ok, code, headers, client} <- request(method, url, req_headers, client_opts),
-         :ok <-
-           header_length_constraint(
-             headers,
-             Keyword.get(opts, :max_body_length, @max_body_length)
-           ) do
+         {:ok, code, headers, client} <-
+           request_with_constraints(method, url, req_headers, client_opts, opts) do
       response(conn, client, url, code, headers, opts)
     else
       {:ok, true} ->
@@ -170,7 +166,8 @@ defmodule Pleroma.ReverseProxy do
       {:ok, code, headers} when code in @valid_resp_codes ->
         {:ok, code, downcase_headers(headers)}
 
-      {:ok, code, _, _} ->
+      {:ok, code, _, client} ->
+        client().close(client)
         {:error, {:invalid_http_response, code}}
 
       {:ok, code, _} ->
@@ -181,6 +178,26 @@ defmodule Pleroma.ReverseProxy do
     end
   end
 
+  defp request_with_constraints(method, url, headers, client_opts, opts) do
+    case request(method, url, headers, client_opts) do
+      {:ok, _code, headers, client} = response ->
+        case header_length_constraint(
+               headers,
+               Keyword.get(opts, :max_body_length, @max_body_length)
+             ) do
+          :ok ->
+            response
+
+          error ->
+            client().close(client)
+            error
+        end
+
+      response ->
+        response
+    end
+  end
+
   defp response(conn, client, url, status, headers, opts) do
     Logger.debug("#{__MODULE__} #{status} #{url} #{inspect(headers)}")
 
@@ -406,7 +423,7 @@ defmodule Pleroma.ReverseProxy do
 
   defp header_length_constraint(_, _), do: :ok
 
-  defp body_size_constraint(size, limit) when is_integer(limit) and limit > 0 and size >= limit do
+  defp body_size_constraint(size, limit) when is_integer(limit) and limit > 0 and size > limit do
     {:error, :body_too_large}
   end
 

lib/pleroma/reverse_proxy/client/tesla.ex

@@ -45,8 +45,8 @@ defmodule Pleroma.ReverseProxy.Client.Tesla do
   @impl true
   @spec stream_body(map()) ::
           {:ok, binary(), map()} | {:error, atom() | String.t()} | :done | no_return()
-  def stream_body(%{pid: pid, fin: true}) do
-    ConnectionPool.release_conn(pid)
+  def stream_body(%{pid: pid, stream: stream, fin: true}) do
+    ConnectionPool.release_stream(pid, stream)
     :done
   end
 
@@ -70,6 +70,10 @@ defmodule Pleroma.ReverseProxy.Client.Tesla do
 
   @impl true
   @spec close(map) :: :ok | no_return()
+  def close(%{pid: pid, stream: stream}) do
+    ConnectionPool.cancel_stream(pid, stream)
+  end
+
   def close(%{pid: pid}) do
     ConnectionPool.release_conn(pid)
   end

lib/pleroma/telemetry/logger.ex

@@ -55,8 +55,8 @@ defmodule Pleroma.Telemetry.Logger do
 
   def handle_event(
         [:pleroma, :connection_pool, :provision_failure],
-        %{opts: [key | _]},
         _,
+        %{key: key},
         _
       ) do
     Logger.debug(fn ->

lib/pleroma/tesla/middleware/connection_pool.ex

@@ -12,25 +12,38 @@ defmodule Pleroma.Tesla.Middleware.ConnectionPool do
   alias Pleroma.Gun.ConnectionPool
 
   @impl Tesla.Middleware
-  def call(%Tesla.Env{url: url, opts: opts} = env, next, _) do
+  def call(%Tesla.Env{opts: opts} = env, next, middleware_opts) do
+    if opts[:adapter][:body_as] == :stream do
+      {:error, :pooled_streaming_requires_chunks}
+    else
+      do_call(env, next, middleware_opts)
+    end
+  end
+
+  defp do_call(%Tesla.Env{url: url, opts: opts} = env, next, _) do
     uri = URI.parse(url)
 
     # Avoid leaking connections when the middleware is called twice
     # with body_as: :chunks. We assume only the middleware can set
     # opts[:adapter][:conn]
-    if opts[:adapter][:conn] do
-      ConnectionPool.release_conn(opts[:adapter][:conn])
-    end
+    cleanup_previous_request(env)
 
     case ConnectionPool.get_conn(uri, opts[:adapter]) do
       {:ok, conn_pid} ->
-        adapter_opts = Keyword.merge(opts[:adapter], conn: conn_pid, close_conn: false)
+        tunnel = ConnectionPool.tunnel_ref(conn_pid)
+
+        adapter_opts =
+          opts[:adapter]
+          |> Keyword.drop([:conn, :stream, :tunnel])
+          |> Keyword.merge(conn: conn_pid, close_conn: false, tunnel: tunnel)
+
         opts = Keyword.put(opts, :adapter, adapter_opts)
         env = %{env | opts: opts}
+        next = use_pool_adapter(next)
 
         case Tesla.run(env, next) do
           {:ok, env} ->
-            unless opts[:adapter][:body_as] == :chunks do
+            unless opts[:adapter][:body_as] == :chunks and chunk_client?(env.body) do
               ConnectionPool.release_conn(conn_pid)
               {_, res} = pop_in(env.opts[:adapter][:conn])
               {:ok, res}
@@ -38,8 +51,11 @@ defmodule Pleroma.Tesla.Middleware.ConnectionPool do
               {:ok, env}
             end
 
+          {:error, {:gun_stream_error, reason}} ->
+            {:error, reason}
+
           err ->
-            ConnectionPool.release_conn(conn_pid)
+            ConnectionPool.discard_conn(conn_pid)
             err
         end
 
@@ -47,4 +63,34 @@ defmodule Pleroma.Tesla.Middleware.ConnectionPool do
         err
     end
   end
+
+  defp cleanup_previous_request(%Tesla.Env{
+         body: %{pid: pid, stream: stream},
+         opts: opts
+       }) do
+    if opts[:adapter][:conn], do: ConnectionPool.cancel_stream(pid, stream)
+  end
+
+  defp cleanup_previous_request(%Tesla.Env{opts: opts}) do
+    conn = opts[:adapter][:conn]
+    stream = opts[:adapter][:stream]
+
+    cond do
+      conn && stream -> ConnectionPool.cancel_stream(conn, stream)
+      conn -> ConnectionPool.release_conn(conn)
+      true -> :ok
+    end
+  end
+
+  defp chunk_client?(%{pid: pid, stream: stream}) when is_pid(pid) and not is_nil(stream),
+    do: true
+
+  defp chunk_client?(_body), do: false
+
+  defp use_pool_adapter(next) do
+    List.update_at(next, -1, fn
+      {Tesla.Adapter.Gun, function, args} -> {Pleroma.HTTP.Adapter.Gun, function, args}
+      adapter -> adapter
+    end)
+  end
 end

lib/pleroma/web/rich_media/helpers.ex

@@ -4,6 +4,8 @@
 
 defmodule Pleroma.Web.RichMedia.Helpers do
   alias Pleroma.Config
+  alias Pleroma.Gun.ConnectionPool
+  alias Pleroma.ReverseProxy.Client.Tesla, as: TeslaClient
 
   require Logger
 
@@ -19,13 +21,31 @@ defmodule Pleroma.Web.RichMedia.Helpers do
   end
 
   defp stream(url) do
-    with {_, {:ok, %Tesla.Env{status: 200, body: stream_body, headers: headers}}} <-
-           {:get, Pleroma.HTTP.get(url, req_headers(), http_options())},
-         {_, :ok} <- {:content_type, check_content_type(headers)},
+    case Pleroma.HTTP.get(url, req_headers(), http_options()) do
+      {:ok, %Tesla.Env{status: 200, body: stream_body, headers: headers}} ->
+        result =
+          try do
+            with {_, :ok} <- {:content_type, check_content_type(headers)},
                  {_, :ok} <- {:content_length, check_content_length(headers)},
                  {:read_stream, {:ok, body}} <- {:read_stream, read_stream(stream_body)} do
               {:ok, body}
             end
+          rescue
+            exception ->
+              cleanup_stream(stream_body, :error)
+              reraise exception, __STACKTRACE__
+          end
+
+        cleanup_stream(stream_body, result)
+        result
+
+      {:ok, %Tesla.Env{body: stream_body}} = response ->
+        cleanup_stream(stream_body, :rejected)
+        {:get, response}
+
+      response ->
+        {:get, response}
+    end
   end
 
   defp head_first(url) do
@@ -95,6 +115,10 @@ defmodule Pleroma.Web.RichMedia.Helpers do
     end
   end
 
+  defp read_stream(%{pid: pid, stream: stream, opts: opts}) do
+    read_chunks(pid, stream, opts, "", 0, Keyword.get(http_options(), :max_body))
+  end
+
   defp read_stream(stream) do
     max_body = Keyword.get(http_options(), :max_body)
 
@@ -117,6 +141,33 @@ defmodule Pleroma.Web.RichMedia.Helpers do
     end
   end
 
+  defp read_chunks(pid, stream, opts, acc, total_bytes, max_body) do
+    case Tesla.Adapter.Gun.read_chunk(pid, stream, opts) do
+      {fin, chunk} when fin in [:fin, :nofin] ->
+        total_bytes = total_bytes + byte_size(chunk)
+
+        cond do
+          total_bytes > max_body ->
+            :error
+
+          fin == :fin ->
+            {:ok, acc <> chunk}
+
+          true ->
+            read_chunks(pid, stream, opts, acc <> chunk, total_bytes, max_body)
+        end
+
+      {:error, _reason} ->
+        :error
+    end
+  end
+
+  defp cleanup_stream(%{pid: pid, stream: stream}, {:ok, _body}),
+    do: ConnectionPool.release_stream(pid, stream)
+
+  defp cleanup_stream(%{pid: _pid} = client, _result), do: TeslaClient.close(client)
+  defp cleanup_stream(_stream, _result), do: :ok
+
   defp http_options do
     [
       pool: :rich_media,

test/pleroma/gun/conn_test.exs

@@ -0,0 +1,144 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Pleroma.Gun.ConnTest do
+  use ExUnit.Case
+
+  import Mox
+
+  alias Pleroma.Gun.Conn
+
+  setup :verify_on_exit!
+
+  setup do
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http} end)
+    stub(Pleroma.GunMock, :await_tunnel_up, fn _, _, _ -> {:ok, :http2} end)
+    stub(Pleroma.GunMock, :close, fn _ -> :ok end)
+
+    :ok
+  end
+
+  test "selects the direct transport from the URI scheme" do
+    expect(Pleroma.GunMock, :open, 2, fn
+      ~c"example.com", 8443, %{transport: :tls, tls_opts: tls_opts} ->
+        assert tls_opts[:verify] == :verify_peer
+        {:ok, spawn_link(fn -> Process.sleep(:infinity) end)}
+
+      ~c"example.com", 443, %{transport: :tcp} ->
+        {:ok, spawn_link(fn -> Process.sleep(:infinity) end)}
+    end)
+
+    assert {:ok, https_conn, :http, nil} = Conn.open(URI.parse("https://example.com:8443"), [])
+    assert {:ok, http_conn, :http, nil} = Conn.open(URI.parse("http://example.com:443"), [])
+
+    stop_conn(https_conn)
+    stop_conn(http_conn)
+  end
+
+  test "passes authentication to an HTTP CONNECT proxy" do
+    conn = spawn_link(fn -> Process.sleep(:infinity) end)
+    stream = make_ref()
+
+    expect(Pleroma.GunMock, :open, fn ~c"proxy.example", 8080, opts ->
+      assert opts.transport == :tcp
+      refute Map.has_key?(opts, :tls_opts)
+      {:ok, conn}
+    end)
+
+    expect(Pleroma.GunMock, :connect, fn ^conn, connect_opts ->
+      assert connect_opts.username == "alice"
+      assert connect_opts.password == "secret"
+      assert connect_opts.transport == :tls
+      assert connect_opts.protocols == [:http2, :http]
+      stream
+    end)
+
+    expect(Pleroma.GunMock, :await, fn ^conn, ^stream -> {:response, :fin, 200, []} end)
+
+    assert {:ok, ^conn, :http2, ^stream} =
+             Conn.open(URI.parse("https://origin.example/inbox"),
+               proxy: {~c"proxy.example", 8080},
+               proxy_auth: {"alice", "secret"}
+             )
+
+    stop_conn(conn)
+  end
+
+  test "uses separate TLS options for a TLS CONNECT proxy" do
+    conn = spawn_link(fn -> Process.sleep(:infinity) end)
+    stream = make_ref()
+
+    expect(Pleroma.GunMock, :open, fn ~c"proxy.example", 8443, opts ->
+      assert opts.transport == :tls
+      assert opts.tls_opts == [verify: :verify_none]
+      {:ok, conn}
+    end)
+
+    expect(Pleroma.GunMock, :connect, fn ^conn, _connect_opts -> stream end)
+    expect(Pleroma.GunMock, :await, fn ^conn, ^stream -> {:response, :fin, 200, []} end)
+
+    assert {:ok, ^conn, :http2, ^stream} =
+             Conn.open(URI.parse("https://origin.example/inbox"),
+               proxy: {~c"proxy.example", 8443},
+               transport: :tls,
+               proxy_tls_opts: [verify: :verify_none]
+             )
+
+    stop_conn(conn)
+  end
+
+  test "passes authentication and remote DNS destination to a SOCKS5 proxy" do
+    conn = spawn_link(fn -> Process.sleep(:infinity) end)
+
+    expect(Pleroma.GunMock, :open, fn ~c"proxy.example", 1080, opts ->
+      assert opts.protocols == [:socks]
+      assert opts.transport == :tcp
+      assert opts.socks_opts.host == ~c"origin.example"
+      assert opts.socks_opts.port == 443
+      assert opts.socks_opts.auth == [{:username_password, "alice", "secret"}]
+      assert opts.socks_opts.transport == :tls
+      assert opts.socks_opts.protocols == [:http2, :http]
+      refute Map.has_key?(opts, :tls_opts)
+      {:ok, conn}
+    end)
+
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :socks} end)
+
+    assert {:ok, ^conn, :http2, nil} =
+             Conn.open(URI.parse("https://origin.example/inbox"),
+               proxy: {:socks5, ~c"proxy.example", 1080},
+               proxy_auth: {"alice", "secret"}
+             )
+
+    stop_conn(conn)
+  end
+
+  test "closes a CONNECT socket when proxy authentication fails" do
+    conn = spawn_link(fn -> Process.sleep(:infinity) end)
+    stream = make_ref()
+
+    expect(Pleroma.GunMock, :open, fn _, _, _ -> {:ok, conn} end)
+    expect(Pleroma.GunMock, :connect, fn ^conn, _ -> stream end)
+    expect(Pleroma.GunMock, :await, fn ^conn, ^stream -> {:response, :nofin, 407, []} end)
+    expect(Pleroma.GunMock, :close, fn ^conn -> Process.exit(conn, :normal) end)
+
+    assert {:error, :proxy_auth_failed} =
+             Conn.open(URI.parse("https://origin.example/inbox"),
+               proxy: {~c"proxy.example", 8080},
+               proxy_auth: {"alice", "wrong"}
+             )
+  end
+
+  test "rejects SOCKS4 before opening a socket" do
+    assert {:error, :socks4_unsupported} =
+             Conn.open(URI.parse("https://origin.example/inbox"),
+               proxy: {:socks4, ~c"proxy.example", 1080}
+             )
+  end
+
+  defp stop_conn(conn) do
+    Process.unlink(conn)
+    Process.exit(conn, :kill)
+  end
+end

test/pleroma/gun/connection_pool_test.exs

@@ -11,8 +11,11 @@ defmodule Pleroma.Gun.ConnectionPoolTest do
 
   defp gun_mock(_) do
     Pleroma.GunMock
-    |> stub(:open, fn _, _, _ -> Task.start_link(fn -> Process.sleep(100) end) end)
+    |> stub(:open, fn _, _, _ -> Task.start_link(fn -> Process.sleep(:infinity) end) end)
     |> stub(:await_up, fn _, _ -> {:ok, :http} end)
+    |> stub(:connect, fn _, _ -> make_ref() end)
+    |> stub(:await, fn _, _ -> {:response, :fin, 200, []} end)
+    |> stub(:await_tunnel_up, fn _, _, _ -> {:ok, :http2} end)
     |> stub(:set_owner, fn _, _ -> :ok end)
 
     :ok
@@ -20,40 +23,213 @@ defmodule Pleroma.Gun.ConnectionPoolTest do
 
   setup :gun_mock
 
-  test "gives the same connection to 2 concurrent requests" do
-    Enum.map(
-      [
-        "http://www.korean-books.com.kp/KBMbooks/en/periodic/pictorial/20200530163914.pdf",
-        "http://www.korean-books.com.kp/KBMbooks/en/periodic/pictorial/20200528183427.pdf"
-      ],
-      fn uri ->
-        uri = URI.parse(uri)
-        task_parent = self()
+  test "opens sibling HTTP/1 connections while an existing connection is in use" do
+    uri = URI.parse("http://www.korean-books.com.kp/document.pdf")
 
-        Task.start_link(fn ->
-          {:ok, conn} = ConnectionPool.get_conn(uri, [])
-          ConnectionPool.release_conn(conn)
-          send(task_parent, conn)
-        end)
+    assert {:ok, first} = ConnectionPool.get_conn(uri, [])
+    assert {:ok, second} = ConnectionPool.get_conn(uri, [])
+    assert first != second
+
+    assert :ok = ConnectionPool.release_conn(first)
+    assert :ok = ConnectionPool.release_conn(second)
   end
+
+  test "multiplexes concurrent HTTP/2 leases on one connection" do
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http2} end)
+    uri = URI.parse("https://h2.example/inbox")
+
+    assert {:ok, conn} = ConnectionPool.get_conn(uri, [])
+    assert {:ok, ^conn} = ConnectionPool.get_conn(uri, [])
+
+    assert :ok = ConnectionPool.release_conn(conn)
+    assert :ok = ConnectionPool.release_conn(conn)
+    assert Process.alive?(conn)
+  end
+
+  test "does not share connections across direct and authenticated proxy routes" do
+    uri = URI.parse("https://routes.example/inbox")
+
+    assert {:ok, direct} = ConnectionPool.get_conn(uri, [])
+
+    assert {:ok, proxied} =
+             ConnectionPool.get_conn(uri,
+               proxy: {~c"proxy.example", 8080},
+               proxy_auth: {"alice", "secret"}
              )
 
-    [pid, pid] =
-      for _ <- 1..2 do
+    assert {:ok, other_credentials} =
+             ConnectionPool.get_conn(uri,
+               proxy: {~c"proxy.example", 8080},
+               proxy_auth: {"bob", "different-secret"}
+             )
+
+    assert Enum.uniq([direct, proxied, other_credentials]) ==
+             [direct, proxied, other_credentials]
+
+    registry_dump = inspect(Registry.select(ConnectionPool, [{{:"$1", :_, :_}, [], [:"$1"]}]))
+    refute registry_dump =~ "secret"
+
+    Enum.each([direct, proxied, other_credentials], &ConnectionPool.release_conn/1)
+  end
+
+  test "coalesces concurrent cold HTTP/2 checkouts behind one connection" do
+    test_pid = self()
+
+    stub(Pleroma.GunMock, :await_up, fn _, _ ->
+      send(test_pid, {:awaiting_connection, self()})
+
       receive do
-          pid -> pid
+        :continue -> {:ok, :http2}
       end
+    end)
+
+    uri = URI.parse("https://cold-h2.example/inbox")
+
+    tasks =
+      for _ <- 1..10 do
+        Task.async(fn -> ConnectionPool.get_conn(uri, []) end)
       end
+
+    assert_receive {:awaiting_connection, worker_pid}
+    refute_receive {:awaiting_connection, _other_worker}, 50
+    send(worker_pid, :continue)
+
+    assert [{:ok, conn}] = tasks |> Enum.map(&Task.await/1) |> Enum.uniq()
+    assert Process.alive?(conn)
+  end
+
+  test "shares a cold connection failure with concurrent waiters" do
+    test_pid = self()
+
+    stub(Pleroma.GunMock, :open, fn _, _, _ ->
+      send(test_pid, {:opening_connection, self()})
+
+      receive do
+        :fail -> {:error, :econnrefused}
+      end
+    end)
+
+    uri = URI.parse("https://unreachable.example/inbox")
+
+    tasks =
+      for _ <- 1..10 do
+        Task.async(fn -> ConnectionPool.get_conn(uri, []) end)
+      end
+
+    assert_receive {:opening_connection, worker_pid}
+    refute_receive {:opening_connection, _other_worker}, 50
+    send(worker_pid, :fail)
+
+    assert Enum.all?(tasks, fn task -> match?({:error, _reason}, Task.await(task)) end)
+    refute_receive {:opening_connection, _other_worker}, 50
+  end
+
+  test "bounds many-host workloads and closes released idle workers" do
+    clear_config([:connections_pool, :max_connections]) do
+      clear_config([:connections_pool, :max_connections], 25)
+      clear_config([:connections_pool, :max_idle_time], 5)
+      restart_worker_supervisor()
+
+      on_exit(&restart_worker_supervisor/0)
+    end
+
+    connections =
+      Enum.map(1..25, fn host_number ->
+        uri = URI.parse("https://host-#{host_number}.example/inbox")
+        assert {:ok, conn} = ConnectionPool.get_conn(uri, [])
+        conn
+      end)
+
+    assert DynamicSupervisor.count_children(Pleroma.Gun.ConnectionPool.WorkerSupervisor).active ==
+             25
+
+    assert {:error, :pool_full} =
+             ConnectionPool.get_conn(URI.parse("https://overflow.example/inbox"), [])
+
+    Enum.each(connections, &ConnectionPool.release_conn/1)
+
+    assert eventually(fn ->
+             DynamicSupervisor.count_children(Pleroma.Gun.ConnectionPool.WorkerSupervisor).active ==
+               0
+           end)
+  end
+
+  test "does not reclaim a worker that became active after idle selection" do
+    uri = URI.parse("https://reclaim-race.example/inbox")
+
+    assert {:ok, conn} = ConnectionPool.get_conn(uri, [])
+    assert :ok = ConnectionPool.release_conn(conn)
+    worker = worker_for_conn(conn)
+
+    assert {:idle, _crf, _last_reference} = GenServer.call(worker, :reclaim_info)
+    assert {:ok, ^conn} = ConnectionPool.get_conn(uri, [])
+    assert :in_use = GenServer.call(worker, :reclaim)
+    assert Process.alive?(worker)
+
+    assert :ok = ConnectionPool.release_conn(conn)
+  end
+
+  test "cancels one HTTP/2 stream without discarding the connection" do
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http2} end)
+    uri = URI.parse("https://cancel-h2.example/inbox")
+    stream = make_ref()
+
+    expect(Pleroma.GunMock, :cancel, fn conn, ^stream when is_pid(conn) -> :ok end)
+
+    assert {:ok, conn} = ConnectionPool.get_conn(uri, [])
+    assert :ok = ConnectionPool.register_stream(conn, stream)
+    assert :ok = ConnectionPool.cancel_stream(conn, stream)
+    assert {:ok, ^conn} = ConnectionPool.get_conn(uri, [])
+    assert :ok = ConnectionPool.release_conn(conn)
+  end
+
+  test "ignores a stale Gun error after cancelling an HTTP/2 stream" do
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http2} end)
+    uri = URI.parse("https://stale-cancel-h2.example/inbox")
+    stream = make_ref()
+
+    expect(Pleroma.GunMock, :cancel, fn conn, ^stream when is_pid(conn) -> :ok end)
+
+    assert {:ok, conn} = ConnectionPool.get_conn(uri, [])
+    worker = worker_for_conn(conn)
+    assert :ok = ConnectionPool.register_stream(conn, stream)
+    assert :ok = ConnectionPool.cancel_stream(conn, stream)
+
+    send(worker, {:gun_error, conn, stream, {:badstate, "The stream cannot be found."}})
+    assert %{clients: %{}} = :sys.get_state(worker)
+    assert Process.alive?(worker)
+
+    assert {:ok, ^conn} = ConnectionPool.get_conn(uri, [])
+    assert :ok = ConnectionPool.release_conn(conn)
+  end
+
+  test "stops an HTTP/1 worker when a client dies with an active stream" do
+    uri = URI.parse("http://dead-stream.example/media")
+    parent = self()
+
+    client =
+      spawn(fn ->
+        {:ok, conn} = ConnectionPool.get_conn(uri, [])
+        :ok = ConnectionPool.register_stream(conn, make_ref())
+        send(parent, {:active_stream, conn})
+        Process.sleep(:infinity)
+      end)
+
+    assert_receive {:active_stream, conn}
+    worker = worker_for_conn(conn)
+    Process.exit(client, :kill)
+
+    assert eventually(fn -> not Process.alive?(worker) end)
   end
 
   test "connection limit is respected with concurrent requests" do
     clear_config([:connections_pool, :max_connections]) do
       clear_config([:connections_pool, :max_connections], 1)
       # The supervisor needs a reboot to apply the new config setting
-      Process.exit(Process.whereis(Pleroma.Gun.ConnectionPool.WorkerSupervisor), :kill)
+      restart_worker_supervisor()
 
       on_exit(fn ->
-        Process.exit(Process.whereis(Pleroma.Gun.ConnectionPool.WorkerSupervisor), :kill)
+        restart_worker_supervisor()
       end)
     end
 
@@ -96,4 +272,41 @@ defmodule Pleroma.Gun.ConnectionPoolTest do
         |> Enum.sort()
     end)
   end
+
+  defp restart_worker_supervisor do
+    supervisor = Process.whereis(Pleroma.Gun.ConnectionPool.WorkerSupervisor)
+    monitor = Process.monitor(supervisor)
+    Process.exit(supervisor, :kill)
+
+    assert_receive {:DOWN, ^monitor, :process, ^supervisor, :killed}
+
+    assert eventually(fn ->
+             case Process.whereis(Pleroma.Gun.ConnectionPool.WorkerSupervisor) do
+               pid when is_pid(pid) -> pid != supervisor and Process.alive?(pid)
+               nil -> false
+             end
+           end)
+  end
+
+  defp worker_for_conn(conn) do
+    [worker] =
+      Registry.select(ConnectionPool, [
+        {{:_, :"$1", {:"$2", :_}}, [{:==, :"$2", conn}], [:"$1"]}
+      ])
+
+    worker
+  end
+
+  defp eventually(predicate, attempts \\ 50)
+
+  defp eventually(_predicate, 0), do: false
+
+  defp eventually(predicate, attempts) do
+    if predicate.() do
+      true
+    else
+      Process.sleep(10)
+      eventually(predicate, attempts - 1)
+    end
+  end
 end

test/pleroma/http/adapter_helper/gun_test.exs

@@ -73,5 +73,11 @@ defmodule Pleroma.HTTP.AdapterHelper.GunTest do
 
       assert opts[:proxy] == {~c"example.com", 4321}
     end
+
+    test "uses managed chunks for streamed responses" do
+      opts = Gun.options([stream: true], URI.parse("https://example.com"))
+
+      assert opts[:body_as] == :chunks
+    end
   end
 end

test/pleroma/http/gun_integration_test.exs

@@ -0,0 +1,671 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Pleroma.HTTP.GunIntegrationTest do
+  use ExUnit.Case, async: false
+  use Pleroma.Tests.Helpers
+
+  alias Pleroma.Tesla.Middleware.ConnectionPool
+
+  setup do
+    clear_config([Pleroma.Gun], Pleroma.Gun.API)
+    clear_config([:connections_pool, :max_idle_time], 1_000)
+
+    origin = start_tls_origin(self())
+    connect_proxy = start_connect_proxy(self())
+    socks_proxy = start_socks5_proxy(self())
+
+    on_exit(fn ->
+      Process.sleep(20)
+      stop_server(socks_proxy)
+      stop_server(connect_proxy)
+      stop_server(origin)
+    end)
+
+    {:ok, origin: origin, connect_proxy: connect_proxy, socks_proxy: socks_proxy}
+  end
+
+  test "requests HTTPS directly on a non-standard port", %{origin: origin} do
+    assert {:ok, %Tesla.Env{status: 200, body: "ok"}} = request(origin.base_url <> "/final")
+  end
+
+  test "follows redirects through an authenticated CONNECT proxy", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    assert {:ok, %Tesla.Env{status: 200, body: "ok"}} =
+             request(origin.base_url <> "/redirect", proxy_opts, true)
+
+    expected = "Basic " <> Base.encode64("alice:secret")
+    assert_receive {:connect_proxy, ^expected, "127.0.0.1", origin_port}
+    assert origin_port == origin.port
+  end
+
+  test "requests through an authenticated CONNECT proxy", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    result = request(origin.base_url <> "/final", proxy_opts)
+
+    expected = "Basic " <> Base.encode64("alice:secret")
+    assert_receive {:connect_proxy, ^expected, "127.0.0.1", origin_port}
+    assert origin_port == origin.port
+    assert worker_protocol(origin.port) == :http
+    assert_receive {:origin_request, "/final"}
+    assert {:ok, %Tesla.Env{status: 200, body: "ok"}} = result
+  end
+
+  test "opens a request-ready authenticated CONNECT tunnel", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    uri = URI.parse(origin.base_url <> "/final")
+
+    opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"},
+      tls_opts: [verify: :verify_none],
+      connect_timeout: 2_000
+    ]
+
+    assert {:ok, conn, :http, tunnel} = Pleroma.Gun.Conn.open(uri, opts)
+    stream = :gun.get(conn, "/final", [], %{tunnel: tunnel})
+    assert {:response, :nofin, 200, _headers} = :gun.await(conn, stream, 2_000)
+    assert {:ok, "ok"} = :gun.await_body(conn, stream, 2_000)
+    :gun.close(conn)
+  end
+
+  test "streams chunks through an authenticated CONNECT tunnel", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    assert {:ok, %Tesla.Env{status: 200, body: client}} =
+             request_chunks(origin.base_url <> "/chunks", proxy_opts)
+
+    assert collect_chunks(client) == "hello"
+    assert :ok = Pleroma.Gun.ConnectionPool.release_conn(client.pid)
+  end
+
+  test "discards an unfinished HTTP/1 redirect stream before following it", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    assert {:ok, %Tesla.Env{status: 200, body: client}} =
+             request_chunks(origin.base_url <> "/redirect_chunks", proxy_opts, true)
+
+    assert collect_chunks(client) == "hello"
+    assert :ok = Pleroma.Gun.ConnectionPool.release_conn(client.pid)
+
+    expected = "Basic " <> Base.encode64("alice:secret")
+    assert_receive {:connect_proxy, ^expected, "127.0.0.1", _port}
+    assert_receive {:connect_proxy, ^expected, "127.0.0.1", _port}
+  end
+
+  test "uploads multipart bodies through an authenticated CONNECT tunnel", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    multipart = Tesla.Multipart.new() |> Tesla.Multipart.add_field("name", "value")
+
+    assert {:ok, %Tesla.Env{status: 200}} =
+             request_multipart(origin.base_url <> "/upload", multipart, proxy_opts)
+
+    assert_receive {:origin_headers, "/upload", headers}
+    assert String.downcase(headers) =~ "content-type: multipart/form-data;"
+    assert_receive {:origin_body, "/upload", body}
+    assert body =~ "name=\"name\""
+    assert body =~ "value"
+  end
+
+  test "uploads streamed bodies through an authenticated CONNECT tunnel", %{
+    origin: origin,
+    connect_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    body = Stream.map(["streamed", "-body"], & &1)
+
+    assert {:ok, %Tesla.Env{status: 200}} =
+             request_stream(origin.base_url <> "/upload", body, proxy_opts)
+
+    assert_receive {:origin_body, "/upload", "streamed-body"}
+  end
+
+  test "discards a partial HTTP/1 upload when body enumeration raises", %{origin: origin} do
+    body =
+      Stream.map(["partial", "raise"], fn
+        "raise" -> raise "upload failed"
+        part -> part
+      end)
+
+    assert_raise RuntimeError, "upload failed", fn ->
+      request_stream(origin.base_url <> "/upload", body, [])
+    end
+
+    assert_receive {:origin_headers, "/upload", _headers}
+    assert worker_protocol(origin.port) == nil
+
+    assert {:ok, %Tesla.Env{status: 200, body: "ok"}} = request(origin.base_url <> "/final")
+  end
+
+  test "requests through authenticated SOCKS5 with proxy-side DNS", %{
+    origin: origin,
+    socks_proxy: proxy
+  } do
+    proxy_opts = [
+      proxy: {:socks5, ~c"127.0.0.1", proxy.port},
+      proxy_auth: {"alice", "secret"}
+    ]
+
+    url = "https://localhost:#{origin.port}/final"
+    assert {:ok, %Tesla.Env{status: 200, body: "ok"}} = request(url, proxy_opts)
+
+    assert_receive {:socks5_proxy, "alice", "secret", "localhost", origin_port}
+    assert origin_port == origin.port
+  end
+
+  defp request(url, proxy_opts \\ [], follow_redirects \\ false) do
+    middleware =
+      if follow_redirects do
+        [Tesla.Middleware.FollowRedirects, ConnectionPool]
+      else
+        [ConnectionPool]
+      end
+
+    adapter_opts =
+      Keyword.merge(
+        [tls_opts: [verify: :verify_none], connect_timeout: 2_000, timeout: 2_000],
+        proxy_opts
+      )
+
+    middleware
+    |> Tesla.client(Tesla.Adapter.Gun)
+    |> Tesla.get(url, opts: [adapter: adapter_opts])
+  end
+
+  defp request_chunks(url, proxy_opts, follow_redirects \\ false) do
+    adapter_opts =
+      Keyword.merge(
+        [
+          body_as: :chunks,
+          tls_opts: [verify: :verify_none],
+          connect_timeout: 2_000,
+          timeout: 2_000
+        ],
+        proxy_opts
+      )
+
+    middleware =
+      if follow_redirects do
+        [Tesla.Middleware.FollowRedirects, ConnectionPool]
+      else
+        [ConnectionPool]
+      end
+
+    middleware
+    |> Tesla.client(Tesla.Adapter.Gun)
+    |> Tesla.get(url, opts: [adapter: adapter_opts])
+  end
+
+  defp request_multipart(url, multipart, proxy_opts) do
+    adapter_opts =
+      Keyword.merge(
+        [tls_opts: [verify: :verify_none], connect_timeout: 2_000, timeout: 2_000],
+        proxy_opts
+      )
+
+    [ConnectionPool]
+    |> Tesla.client(Tesla.Adapter.Gun)
+    |> Tesla.post(url, multipart, opts: [adapter: adapter_opts])
+  end
+
+  defp request_stream(url, body, proxy_opts) do
+    adapter_opts =
+      Keyword.merge(
+        [tls_opts: [verify: :verify_none], connect_timeout: 2_000, timeout: 2_000],
+        proxy_opts
+      )
+
+    [ConnectionPool]
+    |> Tesla.client(Tesla.Adapter.Gun)
+    |> Tesla.post(url, body, opts: [adapter: adapter_opts])
+  end
+
+  defp collect_chunks(client, acc \\ "") do
+    case Tesla.Adapter.Gun.read_chunk(client.pid, client.stream, client.opts) do
+      {:nofin, chunk} -> collect_chunks(client, acc <> chunk)
+      {:fin, chunk} -> acc <> chunk
+    end
+  end
+
+  defp start_tls_origin(parent) do
+    certfile = Path.expand("../../fixtures/server.pem", __DIR__)
+    keyfile = Path.expand("../../fixtures/private_key.pem", __DIR__)
+
+    {:ok, listener} =
+      :ssl.listen(0, [
+        :binary,
+        certfile: certfile,
+        keyfile: keyfile,
+        reuseaddr: true,
+        active: false,
+        packet: :raw,
+        ip: {127, 0, 0, 1}
+      ])
+
+    {:ok, {{127, 0, 0, 1}, port}} = :ssl.sockname(listener)
+    {:ok, acceptor} = Task.start_link(fn -> accept_tls_loop(listener, parent) end)
+
+    %{
+      listener: listener,
+      acceptor: acceptor,
+      transport: :ssl,
+      port: port,
+      base_url: "https://127.0.0.1:#{port}"
+    }
+  end
+
+  defp accept_tls_loop(listener, parent) do
+    case :ssl.transport_accept(listener) do
+      {:ok, tcp_socket} ->
+        case :ssl.handshake(tcp_socket, 2_000) do
+          {:ok, socket} ->
+            pid = spawn(fn -> receive_and_serve_tls(socket, parent) end)
+            :ok = :ssl.controlling_process(socket, pid)
+            send(pid, :serve)
+
+          {:error, _reason} ->
+            :gen_tcp.close(tcp_socket)
+        end
+
+        accept_tls_loop(listener, parent)
+
+      {:error, _reason} ->
+        :ok
+    end
+  end
+
+  defp receive_and_serve_tls(socket, parent) do
+    receive do
+      :serve -> serve_tls(socket, parent)
+    end
+  end
+
+  defp serve_tls(socket, parent) do
+    case recv_headers(:ssl, socket) do
+      {:ok, request} ->
+        {headers, buffered_body} = split_headers(request)
+        path = request_path(headers)
+        send(parent, {:origin_request, path})
+        send(parent, {:origin_headers, path, headers})
+
+        case path do
+          "/redirect" ->
+            send_tls_response(socket, 302, "Found", [{"Location", "/final"}], "")
+
+          "/redirect_chunks" ->
+            send_tls_response(socket, 302, "Found", [{"Location", "/chunks"}], "discard-me")
+
+          "/final" ->
+            send_tls_response(socket, 200, "OK", [], "ok")
+
+          "/chunks" ->
+            send_tls_chunks(socket, ["hel", "lo"])
+
+          "/upload" ->
+            {:ok, body} = recv_request_body(socket, headers, buffered_body)
+            send(parent, {:origin_body, path, body})
+            send_tls_response(socket, 200, "OK", [], "uploaded")
+
+          _ ->
+            send_tls_response(socket, 404, "Not Found", [], "not found")
+        end
+
+        serve_tls(socket, parent)
+
+      {:error, _reason} ->
+        :ssl.close(socket)
+    end
+  end
+
+  defp send_tls_response(socket, status, reason, headers, body) do
+    headers =
+      [
+        {"Content-Length", Integer.to_string(byte_size(body))},
+        {"Connection", "keep-alive"}
+      ] ++ headers
+
+    :ssl.send(socket, [
+      "HTTP/1.1 ",
+      Integer.to_string(status),
+      " ",
+      reason,
+      "\r\n",
+      Enum.map(headers, fn {key, value} -> [key, ": ", value, "\r\n"] end),
+      "\r\n",
+      body
+    ])
+  end
+
+  defp send_tls_chunks(socket, chunks) do
+    :ssl.send(socket, [
+      "HTTP/1.1 200 OK\r\n",
+      "Transfer-Encoding: chunked\r\n",
+      "Connection: keep-alive\r\n\r\n",
+      Enum.map(chunks, fn chunk ->
+        [Integer.to_string(byte_size(chunk), 16), "\r\n", chunk, "\r\n"]
+      end),
+      "0\r\n\r\n"
+    ])
+  end
+
+  defp start_connect_proxy(parent) do
+    start_tcp_server(fn socket -> serve_connect_proxy(socket, parent) end)
+  end
+
+  defp serve_connect_proxy(socket, parent) do
+    with {:ok, headers} <- recv_headers(:gen_tcp, socket),
+         {:ok, host, port} <- parse_connect(headers),
+         auth when is_binary(auth) <- header(headers, "proxy-authorization"),
+         {:ok, upstream} <- tcp_connect(host, port) do
+      send(parent, {:connect_proxy, auth, host, port})
+      :ok = :gen_tcp.send(socket, "HTTP/1.1 200 Connection established\r\n\r\n")
+      tunnel(socket, upstream)
+    else
+      _ ->
+        :gen_tcp.send(
+          socket,
+          "HTTP/1.1 407 Proxy Authentication Required\r\nContent-Length: 0\r\n\r\n"
+        )
+
+        :gen_tcp.close(socket)
+    end
+  end
+
+  defp start_socks5_proxy(parent) do
+    start_tcp_server(fn socket -> serve_socks5_proxy(socket, parent) end)
+  end
+
+  defp serve_socks5_proxy(socket, parent) do
+    with {:ok, <<5, method_count>>} <- :gen_tcp.recv(socket, 2, 2_000),
+         {:ok, methods} <- :gen_tcp.recv(socket, method_count, 2_000),
+         true <- :binary.match(methods, <<2>>) != :nomatch,
+         :ok <- :gen_tcp.send(socket, <<5, 2>>),
+         {:ok, username, password} <- recv_socks5_auth(socket),
+         :ok <- :gen_tcp.send(socket, <<1, 0>>),
+         {:ok, host, port} <- recv_socks5_destination(socket),
+         {:ok, upstream} <- tcp_connect(host, port),
+         :ok <- :gen_tcp.send(socket, <<5, 0, 0, 1, 0, 0, 0, 0, 0, 0>>) do
+      send(parent, {:socks5_proxy, username, password, host, port})
+      tunnel(socket, upstream)
+    else
+      _ ->
+        :gen_tcp.close(socket)
+    end
+  end
+
+  defp recv_socks5_auth(socket) do
+    with {:ok, <<1, username_length>>} <- :gen_tcp.recv(socket, 2, 2_000),
+         {:ok, username} <- :gen_tcp.recv(socket, username_length, 2_000),
+         {:ok, <<password_length>>} <- :gen_tcp.recv(socket, 1, 2_000),
+         {:ok, password} <- :gen_tcp.recv(socket, password_length, 2_000) do
+      {:ok, username, password}
+    end
+  end
+
+  defp recv_socks5_destination(socket) do
+    with {:ok, <<5, 1, 0, address_type>>} <- :gen_tcp.recv(socket, 4, 2_000),
+         {:ok, host} <- recv_socks5_host(socket, address_type),
+         {:ok, <<port::16>>} <- :gen_tcp.recv(socket, 2, 2_000) do
+      {:ok, host, port}
+    end
+  end
+
+  defp recv_socks5_host(socket, 1) do
+    with {:ok, <<a, b, c, d>>} <- :gen_tcp.recv(socket, 4, 2_000) do
+      {:ok, Enum.join([a, b, c, d], ".")}
+    end
+  end
+
+  defp recv_socks5_host(socket, 3) do
+    with {:ok, <<length>>} <- :gen_tcp.recv(socket, 1, 2_000),
+         {:ok, host} <- :gen_tcp.recv(socket, length, 2_000) do
+      {:ok, host}
+    end
+  end
+
+  defp start_tcp_server(handler) do
+    {:ok, listener} =
+      :gen_tcp.listen(0, [
+        :binary,
+        active: false,
+        packet: :raw,
+        reuseaddr: true,
+        ip: {127, 0, 0, 1}
+      ])
+
+    {:ok, {{127, 0, 0, 1}, port}} = :inet.sockname(listener)
+    {:ok, acceptor} = Task.start_link(fn -> accept_tcp_loop(listener, handler) end)
+    %{listener: listener, acceptor: acceptor, transport: :gen_tcp, port: port}
+  end
+
+  defp accept_tcp_loop(listener, handler) do
+    case :gen_tcp.accept(listener) do
+      {:ok, socket} ->
+        pid = spawn(fn -> receive_and_serve_tcp(socket, handler) end)
+        :ok = :gen_tcp.controlling_process(socket, pid)
+        send(pid, :serve)
+        accept_tcp_loop(listener, handler)
+
+      {:error, _reason} ->
+        :ok
+    end
+  end
+
+  defp receive_and_serve_tcp(socket, handler) do
+    receive do
+      :serve -> handler.(socket)
+    end
+  end
+
+  defp recv_headers(transport, socket, acc \\ <<>>) do
+    case transport.recv(socket, 0, 2_000) do
+      {:ok, data} ->
+        acc = acc <> data
+
+        if :binary.match(acc, "\r\n\r\n") == :nomatch do
+          recv_headers(transport, socket, acc)
+        else
+          {:ok, acc}
+        end
+
+      {:error, _reason} = error ->
+        error
+    end
+  end
+
+  defp request_path(headers) do
+    headers
+    |> String.split("\r\n", parts: 2)
+    |> hd()
+    |> String.split(" ")
+    |> Enum.at(1)
+  end
+
+  defp split_headers(request) do
+    {index, 4} = :binary.match(request, "\r\n\r\n")
+    split_at = index + 4
+    <<headers::binary-size(split_at), body::binary>> = request
+    {headers, body}
+  end
+
+  defp recv_request_body(socket, headers, buffered_body) do
+    cond do
+      String.contains?(String.downcase(headers), "transfer-encoding: chunked") ->
+        recv_chunked_body(socket, buffered_body, "")
+
+      content_length = header(headers, "content-length") ->
+        recv_exact(socket, buffered_body, String.to_integer(content_length))
+
+      true ->
+        {:ok, buffered_body}
+    end
+  end
+
+  defp recv_chunked_body(socket, buffer, acc) do
+    with {:ok, line, buffer} <- recv_line(socket, buffer),
+         {size, ""} <- Integer.parse(line, 16) do
+      if size == 0 do
+        {:ok, acc}
+      else
+        with {:ok, chunk_and_crlf, buffer} <- take_bytes(socket, buffer, size + 2),
+             <<chunk::binary-size(size), "\r\n">> <- chunk_and_crlf do
+          recv_chunked_body(socket, buffer, acc <> chunk)
+        end
+      end
+    end
+  end
+
+  defp recv_line(socket, buffer) do
+    case :binary.match(buffer, "\r\n") do
+      {index, 2} ->
+        <<line::binary-size(index), "\r\n", rest::binary>> = buffer
+        {:ok, line, rest}
+
+      :nomatch ->
+        with {:ok, data} <- :ssl.recv(socket, 0, 2_000) do
+          recv_line(socket, buffer <> data)
+        end
+    end
+  end
+
+  defp take_bytes(_socket, buffer, size) when byte_size(buffer) >= size do
+    <<bytes::binary-size(size), rest::binary>> = buffer
+    {:ok, bytes, rest}
+  end
+
+  defp take_bytes(socket, buffer, size) do
+    with {:ok, data} <- :ssl.recv(socket, 0, 2_000) do
+      take_bytes(socket, buffer <> data, size)
+    end
+  end
+
+  defp recv_exact(socket, buffer, size) do
+    with {:ok, body, _rest} <- take_bytes(socket, buffer, size) do
+      {:ok, body}
+    end
+  end
+
+  defp parse_connect(headers) do
+    case headers |> String.split("\r\n", parts: 2) |> hd() |> String.split(" ") do
+      ["CONNECT", authority, _protocol] ->
+        case String.split(authority, ":", parts: 2) do
+          [host, port] -> {:ok, host, String.to_integer(port)}
+          _ -> {:error, :invalid_authority}
+        end
+
+      _ ->
+        {:error, :invalid_connect}
+    end
+  end
+
+  defp header(headers, wanted_name) do
+    headers
+    |> String.split("\r\n")
+    |> Enum.find_value(fn line ->
+      case String.split(line, ":", parts: 2) do
+        [name, value] ->
+          if String.downcase(name) == wanted_name, do: String.trim(value)
+
+        _ ->
+          nil
+      end
+    end)
+  end
+
+  defp tcp_connect(host, port) do
+    :gen_tcp.connect(
+      String.to_charlist(host),
+      port,
+      [:binary, active: false, packet: :raw],
+      2_000
+    )
+  end
+
+  defp worker_protocol(port) do
+    prefix = "https:127.0.0.1:#{port}:"
+
+    Pleroma.Gun.ConnectionPool
+    |> Registry.select([{{:"$1", :"$2", :_}, [], [{{:"$1", :"$2"}}]}])
+    |> Enum.find_value(fn {key, worker} ->
+      if String.starts_with?(key, prefix), do: :sys.get_state(worker).protocol
+    end)
+  end
+
+  defp tunnel(client, upstream) do
+    :ok = :inet.setopts(client, active: :once)
+    :ok = :inet.setopts(upstream, active: :once)
+    tunnel_loop(client, upstream)
+  end
+
+  defp tunnel_loop(client, upstream) do
+    receive do
+      {:tcp, ^client, data} ->
+        :ok = :gen_tcp.send(upstream, data)
+        :ok = :inet.setopts(client, active: :once)
+        tunnel_loop(client, upstream)
+
+      {:tcp, ^upstream, data} ->
+        :ok = :gen_tcp.send(client, data)
+        :ok = :inet.setopts(upstream, active: :once)
+        tunnel_loop(client, upstream)
+
+      {:tcp_closed, _socket} ->
+        :gen_tcp.close(client)
+        :gen_tcp.close(upstream)
+
+      {:tcp_error, _socket, _reason} ->
+        :gen_tcp.close(client)
+        :gen_tcp.close(upstream)
+    after
+      10_000 ->
+        :gen_tcp.close(client)
+        :gen_tcp.close(upstream)
+    end
+  end
+
+  defp stop_server(%{listener: listener, acceptor: acceptor, transport: transport}) do
+    transport.close(listener)
+
+    if Process.alive?(acceptor), do: Process.exit(acceptor, :normal)
+  end
+end

test/pleroma/reverse_proxy/client/tesla_test.exs

@@ -0,0 +1,22 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Pleroma.ReverseProxy.Client.TeslaTest do
+  use ExUnit.Case
+
+  import Mox
+
+  alias Pleroma.ReverseProxy.Client.Tesla
+
+  setup :verify_on_exit!
+
+  test "cancels an unfinished stream before releasing its connection" do
+    conn = self()
+    stream = [make_ref(), make_ref()]
+
+    expect(Pleroma.GunMock, :cancel, fn ^conn, ^stream -> :ok end)
+
+    assert :ok = Tesla.close(%{pid: conn, stream: stream})
+  end
+end

test/pleroma/reverse_proxy_test.exs

@@ -115,6 +115,7 @@ defmodule Pleroma.ReverseProxyTest do
   describe "max_body" do
     test "length returns error if content-length more than option", %{conn: conn} do
       request_mock(0)
+      expect(ClientMock, :close, fn _ -> :ok end)
 
       assert capture_log(fn ->
                ReverseProxy.call(conn, "/huge-file", max_body_length: 4)
@@ -128,14 +129,31 @@ defmodule Pleroma.ReverseProxyTest do
              end) == ""
     end
 
+    test "closes streamed responses with invalid status", %{conn: conn} do
+      ClientMock
+      |> expect(:request, fn :get, "/invalid-status", _, _, _ ->
+        {:ok, 404, [], %{url: "/invalid-status"}}
+      end)
+      |> expect(:close, fn %{url: "/invalid-status"} -> :ok end)
+
+      ReverseProxy.call(conn, "/invalid-status")
+    end
+
     test "max_body_length returns error if streaming body more than that option", %{conn: conn} do
       stream_mock(3, true)
 
       assert capture_log(fn ->
-               ReverseProxy.call(conn, "/stream-bytes/50", max_body_length: 30)
+               ReverseProxy.call(conn, "/stream-bytes/50", max_body_length: 29)
              end) =~
                "Elixir.Pleroma.ReverseProxy request to /stream-bytes/50 failed while reading/chunking: :body_too_large"
     end
+
+    test "max_body_length accepts a body exactly at the limit", %{conn: conn} do
+      stream_mock(4)
+
+      conn = ReverseProxy.call(conn, "/stream-bytes/30", max_body_length: 30)
+      assert byte_size(conn.resp_body) == 30
+    end
   end
 
   describe "HEAD requests" do
@@ -200,7 +218,10 @@ defmodule Pleroma.ReverseProxyTest do
 
     test "204", %{conn: conn} do
       url = "/status/204"
-      expect(ClientMock, :request, fn :get, _url, _, _, _ -> {:ok, 204, [], %{}} end)
+
+      ClientMock
+      |> expect(:request, fn :get, _url, _, _, _ -> {:ok, 204, [], %{}} end)
+      |> expect(:close, fn %{} -> :ok end)
 
       capture_log(fn ->
         conn = ReverseProxy.call(conn, url)

test/pleroma/tesla/middleware/connection_pool_test.exs

@@ -0,0 +1,99 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Pleroma.Tesla.Middleware.ConnectionPoolTest do
+  use Pleroma.DataCase
+
+  import Mox
+
+  alias Pleroma.Tesla.Middleware.ConnectionPool
+
+  defmodule FailingAdapter do
+    @behaviour Tesla.Adapter
+
+    @impl Tesla.Adapter
+    def call(_env, _opts), do: {:error, :request_failed}
+  end
+
+  defmodule EmptyAdapter do
+    @behaviour Tesla.Adapter
+
+    @impl Tesla.Adapter
+    def call(env, _opts), do: {:ok, %{env | status: 204, body: ""}}
+  end
+
+  test "discards the connection when the adapter cannot finish a request" do
+    test_pid = self()
+
+    stub(Pleroma.GunMock, :open, fn _, _, _ ->
+      {:ok, conn} = Task.start_link(fn -> Process.sleep(:infinity) end)
+      send(test_pid, {:gun_conn, conn})
+      {:ok, conn}
+    end)
+
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http} end)
+
+    client = Tesla.client([ConnectionPool], FailingAdapter)
+
+    assert {:error, :request_failed} =
+             Tesla.get(client, "http://discard.example/", opts: [adapter: []])
+
+    assert_receive {:gun_conn, conn}
+
+    assert eventually(fn -> not Process.alive?(conn) end)
+  end
+
+  test "releases a chunks lease when the response has no stream" do
+    clear_config([:connections_pool, :max_idle_time], 5)
+    test_pid = self()
+
+    stub(Pleroma.GunMock, :open, fn _, _, _ ->
+      {:ok, conn} = Task.start_link(fn -> Process.sleep(:infinity) end)
+      send(test_pid, {:gun_conn, conn})
+      {:ok, conn}
+    end)
+
+    stub(Pleroma.GunMock, :await_up, fn _, _ -> {:ok, :http} end)
+
+    client = Tesla.client([ConnectionPool], EmptyAdapter)
+
+    assert {:ok, %Tesla.Env{status: 204, body: ""}} =
+             Tesla.get(client, "http://empty-chunks.example/",
+               opts: [adapter: [body_as: :chunks]]
+             )
+
+    assert_receive {:gun_conn, conn}
+    worker = worker_for_conn(conn)
+    assert eventually(fn -> not Process.alive?(worker) end)
+  end
+
+  test "rejects unmanaged pooled streams" do
+    client = Tesla.client([ConnectionPool], EmptyAdapter)
+
+    assert {:error, :pooled_streaming_requires_chunks} =
+             Tesla.get(client, "http://stream.example/", opts: [adapter: [body_as: :stream]])
+  end
+
+  defp worker_for_conn(conn) do
+    [worker] =
+      Registry.select(Pleroma.Gun.ConnectionPool, [
+        {{:_, :"$1", {:"$2", :_}}, [{:==, :"$2", conn}], [:"$1"]}
+      ])
+
+    worker
+  end
+
+  defp eventually(predicate, attempts \\ 50)
+
+  defp eventually(_predicate, 0), do: false
+
+  defp eventually(predicate, attempts) do
+    if predicate.() do
+      true
+    else
+      Process.sleep(10)
+      eventually(predicate, attempts - 1)
+    end
+  end
+end