diff --git a/lib/plug/conn.ex b/lib/plug/conn.ex index 5780ca71..d7177966 100644 --- a/lib/plug/conn.ex +++ b/lib/plug/conn.ex @@ -294,7 +294,7 @@ defmodule Plug.Conn do alias Plug.Conn @epoch {{1970, 1, 1}, {0, 0, 0}} @already_sent {:plug_conn, :sent} - @unsent [:unset, :set, :set_chunked, :set_file] + @unsent [:unset, :set, :set_upgrade, :set_chunked, :set_file] @doc """ Assigns a value to a key in the connection. @@ -1471,10 +1471,11 @@ defmodule Plug.Conn do @spec upgrade_adapter(t, atom, term) :: t def upgrade_adapter(%Conn{adapter: {adapter, payload}, state: state} = conn, protocol, args) when state in @unsent do + conn = run_before_send(%{conn | status: 101}, :set_upgrade) + case adapter.upgrade(payload, protocol, args) do {:ok, payload} -> - conn = run_before_send(conn, :upgraded) - %{conn | adapter: {adapter, payload}} + %{conn | state: :upgraded, adapter: {adapter, payload}} {:error, :not_supported} -> raise ArgumentError, "upgrade to #{protocol} not supported by #{inspect(adapter)}" diff --git a/test/plug/conn_test.exs b/test/plug/conn_test.exs index d5080944..95af3c85 100644 --- a/test/plug/conn_test.exs +++ b/test/plug/conn_test.exs @@ -510,9 +510,19 @@ defmodule Plug.ConnTest do test "upgrade_adapter/3 runs before_send callbacks" do conn = conn(:get, "/foo") - |> register_before_send(&assign(&1, :ran_before_send, true)) + |> register_before_send(fn conn -> + send(self(), {:state, conn.state}) + assert conn.status == 101 + + conn + |> put_resp_header("x-test", "UPGRADE") + |> assign(:ran_before_send, true) + end) |> upgrade_adapter(:supported, opt: :supported) + assert_receive {:state, :set_upgrade} + + assert {"x-test", "UPGRADE"} in conn.resp_headers assert conn.assigns[:ran_before_send] == true end