Add support for protocol upgrades

lib/plug/adapters/test/conn.ex

@@ -119,6 +119,16 @@ defmodule Plug.Adapters.Test.Conn do
     :ok
   end
 
+  def upgrade(%{owner: owner, ref: ref}, :not_supported = protocol, opts) do
+    send(owner, {ref, :upgrade, {protocol, opts}})
+    {:error, :not_supported}
+  end
+
+  def upgrade(%{owner: owner, ref: ref} = state, protocol, opts) do
+    send(owner, {ref, :upgrade, {protocol, opts}})
+    {:ok, state}
+  end
+
   def push(%{owner: owner, ref: ref}, path, headers) do
     send(owner, {ref, :push, {path, headers}})
     :ok

lib/plug/conn.ex

@@ -108,8 +108,8 @@ defmodule Plug.Conn do
   The connection state is used to track the connection lifecycle. It starts as
   `:unset` but is changed to `:set` (via `resp/3`) or `:set_chunked`
   (used only for `before_send` callbacks by `send_chunked/2`) or `:file`
-  (when invoked via `send_file/3`). Its final result is `:sent`, `:file` or
-  `:chunked` depending on the response model.
+  (when invoked via `send_file/3`). Its final result is `:sent`, `:file`, `:chunked` 
+  or `:upgraded` depending on the response model.
 
   ## Private fields
 
@@ -150,6 +150,17 @@ defmodule Plug.Conn do
   Even though 404 has been overridden, the `:not_found` atom can still be used
   to set the status to 404 as well as the new atom `:actually_this_was_found`
   inflected from the reason phrase "Actually This Was Found".
+
+  ## Protocol Upgrades
+
+  Plug provides basic support for protocol upgrades via the `upgrade_adapter/3`
+  function to facilitate connection upgrades to protocols such as WebSockets.
+  As the name suggests, this functionality is adapter-dependent and the 
+  functionality & requirements of a given upgrade require explicit coordination
+  between a Plug application & the underlying adapter. Plug provides upgrade 
+  related functionality only to the extent necessary to allow a Plug application 
+  to request protocol upgrades from the underlying adapter. See the documentation
+  for `upgrade_adapter/3` for details.
   """
 
   @type adapter :: {module, term}
@@ -172,7 +183,7 @@ defmodule Plug.Conn do
   @type scheme :: :http | :https
   @type secret_key_base :: binary | nil
   @type segments :: [binary]
-  @type state :: :unset | :set | :set_chunked | :set_file | :file | :chunked | :sent
+  @type state :: :unset | :set | :set_chunked | :set_file | :file | :chunked | :sent | :upgraded
   @type status :: atom | int_status
 
   @type t :: %__MODULE__{
@@ -390,7 +401,7 @@ defmodule Plug.Conn do
   atoms is available in `Plug.Conn.Status`.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   ## Examples
 
@@ -399,7 +410,10 @@ defmodule Plug.Conn do
 
   """
   @spec put_status(t, status) :: t
-  def put_status(%Conn{state: :sent}, _status), do: raise(AlreadySentError)
+  def put_status(%Conn{state: state}, _status) when state not in @unsent do
+    raise AlreadySentError
+  end
+
   def put_status(%Conn{} = conn, nil), do: %{conn | status: nil}
   def put_status(%Conn{} = conn, status), do: %{conn | status: Plug.Conn.Status.code(status)}
 
@@ -408,7 +422,7 @@ defmodule Plug.Conn do
 
   It expects the connection state to be `:set`, otherwise raises an
   `ArgumentError` for `:unset` connections or a `Plug.Conn.AlreadySentError` for
-  already `:sent` connections.
+  already `:sent`, `:chunked` or `:upgraded` connections.
 
   At the end sets the connection state to `:sent`.
 
@@ -451,7 +465,7 @@ defmodule Plug.Conn do
   If available, the file is sent directly over the socket using
   the operating system `sendfile` operation.
 
-  It expects a connection that has not been `:sent` yet and sets its
+  It expects a connection that has not been `:sent`, `:chunked` or `:upgraded` yet and sets its
   state to `:file` afterwards. Otherwise raises `Plug.Conn.AlreadySentError`.
 
   ## Examples
@@ -494,7 +508,7 @@ defmodule Plug.Conn do
   @doc """
   Sends the response headers as a chunked response.
 
-  It expects a connection that has not been `:sent` yet and sets its
+  It expects a connection that has not been `:sent` or `:upgraded` yet and sets its
   state to `:chunked` afterwards. Otherwise, raises `Plug.Conn.AlreadySentError`.
   After `send_chunked/2` is called, chunks can be sent to the client via
   the `chunk/2` function.
@@ -590,7 +604,7 @@ defmodule Plug.Conn do
   Sets the response to the given `status` and `body`.
 
   It sets the connection state to `:set` (if not already `:set`)
-  and raises `Plug.Conn.AlreadySentError` if it was already `:sent`.
+  and raises `Plug.Conn.AlreadySentError` if it was already `:sent`, `:chunked` or `:upgraded`.
 
   If you also want to send the response, use `send_resp/1` after this
   or use `send_resp/3`.
@@ -674,7 +688,7 @@ defmodule Plug.Conn do
   headers that aren't lowercase will raise a `Plug.Conn.InvalidHeaderError`.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   ## Examples
 
@@ -684,11 +698,7 @@ defmodule Plug.Conn do
   @spec prepend_req_headers(t, headers) :: t
   def prepend_req_headers(conn, headers)
 
-  def prepend_req_headers(%Conn{state: :sent}, _headers) do
-    raise AlreadySentError
-  end
-
-  def prepend_req_headers(%Conn{state: :chunked}, _headers) do
+  def prepend_req_headers(%Conn{state: state}, _headers) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -723,11 +733,7 @@ defmodule Plug.Conn do
   @spec merge_req_headers(t, Enum.t()) :: t
   def merge_req_headers(conn, headers)
 
-  def merge_req_headers(%Conn{state: :sent}, _headers) do
-    raise AlreadySentError
-  end
-
-  def merge_req_headers(%Conn{state: :chunked}, _headers) do
+  def merge_req_headers(%Conn{state: state}, _headers) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -762,7 +768,7 @@ defmodule Plug.Conn do
   headers that aren't lowercase will raise a `Plug.Conn.InvalidHeaderError`.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   ## Examples
 
@@ -772,7 +778,7 @@ defmodule Plug.Conn do
   @spec put_req_header(t, binary, binary) :: t
   def put_req_header(conn, key, value)
 
-  def put_req_header(%Conn{state: :sent}, _key, _value) do
+  def put_req_header(%Conn{state: state}, _key, _value) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -786,7 +792,7 @@ defmodule Plug.Conn do
   Deletes a request header if present.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   ## Examples
 
@@ -796,11 +802,7 @@ defmodule Plug.Conn do
   @spec delete_req_header(t, binary) :: t
   def delete_req_header(conn, key)
 
-  def delete_req_header(%Conn{state: :sent}, _key) do
-    raise AlreadySentError
-  end
-
-  def delete_req_header(%Conn{state: :chunked}, _key) do
+  def delete_req_header(%Conn{state: state}, _key) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -814,7 +816,7 @@ defmodule Plug.Conn do
   value.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   Only the first value of the header `key` is updated if present.
 
@@ -831,11 +833,7 @@ defmodule Plug.Conn do
   @spec update_req_header(t, binary, binary, (binary -> binary)) :: t
   def update_req_header(conn, key, initial, fun)
 
-  def update_req_header(%Conn{state: :sent}, _key, _initial, _fun) do
-    raise AlreadySentError
-  end
-
-  def update_req_header(%Conn{state: :chunked}, _key, _initial, _fun) do
+  def update_req_header(%Conn{state: state}, _key, _initial, _fun) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -875,7 +873,7 @@ defmodule Plug.Conn do
   headers that aren't lowercase will raise a `Plug.Conn.InvalidHeaderError`.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   Raises a `Plug.Conn.InvalidHeaderError` if the header value contains control
   feed (`\r`) or newline (`\n`) characters.
@@ -886,11 +884,7 @@ defmodule Plug.Conn do
 
   """
   @spec put_resp_header(t, binary, binary) :: t
-  def put_resp_header(%Conn{state: :sent}, _key, _value) do
-    raise AlreadySentError
-  end
-
-  def put_resp_header(%Conn{state: :chunked}, _key, _value) do
+  def put_resp_header(%Conn{state: state}, _key, _value) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -916,7 +910,7 @@ defmodule Plug.Conn do
   headers that aren't lowercase will raise a `Plug.Conn.InvalidHeaderError`.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   Raises a `Plug.Conn.InvalidHeaderError` if the header value contains control
   feed (`\r`) or newline (`\n`) characters.
@@ -929,11 +923,7 @@ defmodule Plug.Conn do
   @spec prepend_resp_headers(t, headers) :: t
   def prepend_resp_headers(conn, headers)
 
-  def prepend_resp_headers(%Conn{state: :sent}, _headers) do
-    raise AlreadySentError
-  end
-
-  def prepend_resp_headers(%Conn{state: :chunked}, _headers) do
+  def prepend_resp_headers(%Conn{state: state}, _headers) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -965,11 +955,7 @@ defmodule Plug.Conn do
   @spec merge_resp_headers(t, Enum.t()) :: t
   def merge_resp_headers(conn, headers)
 
-  def merge_resp_headers(%Conn{state: :sent}, _headers) do
-    raise AlreadySentError
-  end
-
-  def merge_resp_headers(%Conn{state: :chunked}, _headers) do
+  def merge_resp_headers(%Conn{state: state}, _headers) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -993,19 +979,15 @@ defmodule Plug.Conn do
   Deletes a response header if present.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   ## Examples
 
       Plug.Conn.delete_resp_header(conn, "content-type")
 
   """
   @spec delete_resp_header(t, binary) :: t
-  def delete_resp_header(%Conn{state: :sent}, _key) do
-    raise AlreadySentError
-  end
-
-  def delete_resp_header(%Conn{state: :chunked}, _key) do
+  def delete_resp_header(%Conn{state: state}, _key) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -1019,7 +1001,7 @@ defmodule Plug.Conn do
   value.
 
   Raises a `Plug.Conn.AlreadySentError` if the connection has already been
-  `:sent` or `:chunked`.
+  `:sent`, `:chunked` or `:upgraded`.
 
   Only the first value of the header `key` is updated if present.
 
@@ -1036,11 +1018,7 @@ defmodule Plug.Conn do
   @spec update_resp_header(t, binary, binary, (binary -> binary)) :: t
   def update_resp_header(conn, key, initial, fun)
 
-  def update_resp_header(%Conn{state: :sent}, _key, _initial, _fun) do
-    raise AlreadySentError
-  end
-
-  def update_resp_header(%Conn{state: :chunked}, _key, _initial, _fun) do
+  def update_resp_header(%Conn{state: state}, _key, _initial, _fun) when state not in @unsent do
     raise AlreadySentError
   end
 
@@ -1388,6 +1366,46 @@ defmodule Plug.Conn do
   defp adapter_inform(%Conn{adapter: {adapter, payload}}, status, headers),
     do: adapter.inform(payload, status, headers)
 
+  @doc """
+  Request a protocol upgrade from the underlying adapter.
+
+  The precise semantics of an upgrade are deliberately left unspecified here in order to
+  support arbitrary upgrades, even to protocols which may not exist today. The primary intent of
+  this function is solely to allow an application to issue an upgrade request, not to manage how
+  a given protocol upgrade takes place or what APIs the application must support in order to serve
+  this updated protocol. For details in this regard, consult the documentation of the underlying
+  adapter (such a Plug.Cowboy or Bandit).
+
+  Takes an argument describing the requested upgrade (for example, `:websocket`), and an argument
+  which contains arbitrary data which the underlying adapter is expected to interpret in the
+  context of the requested upgrade.
+
+  If the upgrade is accepted by the adapter, the returned `Plug.Conn` will have a `state` of
+  `:upgraded`. This state is considered equivalently to a 'sent' state, and is subject to the same
+  limitation on subsequent mutating operations. Note that there is no guarantee or expectation
+  that the actual upgrade process is undertaken within this function; it is entirely possible that
+  the server will only do the actual upgrade later in the connection lifecycle.
+
+  If the adapter does not support the requested upgrade then this is a noop and the returned 
+  `Plug.Conn` will be unchanged. The application can detect this and operate on the conn as it
+  normally would in order to indicate an upgrade failure to the client.
+  """
+  @spec upgrade_adapter(t, atom, term) :: t
+  def upgrade_adapter(%Conn{adapter: {adapter, payload}, state: state} = conn, protocol, args)
+      when state in @unsent do
+    case adapter.upgrade(payload, protocol, args) do
+      {:ok, payload} ->
+        %{conn | adapter: {adapter, payload}, state: :upgraded}
+
+      {:error, :not_supported} ->
+        raise ArgumentError, "upgrade to #{protocol} not supported by #{inspect(adapter)}"
+    end
+  end
+
+  def upgrade_adapter(_conn, _protocol, _args) do
+    raise AlreadySentError
+  end
+
   @doc """
   Pushes a resource to the client.
 
@@ -1847,8 +1865,10 @@ defmodule Plug.Conn do
     validate_header_value!("set-cookie", cookie)
   end
 
-  defp update_cookies(%Conn{state: :sent}, _fun), do: raise(AlreadySentError)
-  defp update_cookies(%Conn{state: :chunked}, _fun), do: raise(AlreadySentError)
+  defp update_cookies(%Conn{state: state}, _fun) when state not in @unsent do
+    raise AlreadySentError
+  end
+
   defp update_cookies(%Conn{cookies: %Unfetched{}} = conn, _fun), do: conn
   defp update_cookies(%Conn{cookies: cookies} = conn, fun), do: %{conn | cookies: fun.(cookies)}
 

lib/plug/conn/adapter.ex

@@ -133,6 +133,20 @@ defmodule Plug.Conn.Adapter do
   @callback inform(payload, status :: Conn.status(), headers :: Keyword.t()) ::
               :ok | {:error, term}
 
+  @doc """
+  Attempt to upgrade the connection with the client.
+
+  If the adapter does not support the indicated upgrade, then `{:error, :not_supported}` should be
+  be returned.
+
+  If the adapter supports the indicated upgrade but is unable to proceed with it (due to
+  a negotiation error, invalid opts being passed to this function, or some other reason), then an
+  arbitrary error may be returned. Note that an adapter does not need to process the actual
+  upgrade within this function; it is a wholly supported failure mode for an adapter to attempt
+  the upgrade process later in the connection lifecycle and fail at that point.
+  """
+  @callback upgrade(payload, protocol :: atom, opts :: term) :: {:ok, payload} | {:error, term}
+
   @doc """
   Returns peer information such as the address, port and ssl cert.
   """