Allow returning custom error code from socket connect

Previously there were only two options, `{:ok, socket}` and `:error`. The
error always returned `403 Forbidden` to the user.

This is quite limiting. In our use case, we do rate limiting in the
connect function. E.g. if one of the phoenix servers goes down or
restarts, then we cannot immediately handle tens of thousands of users
connecting at the same time. Instead we let in some users, and let
others reconnect later. Also if we're under a heavy load then we want to
limit users connecting.

"Forbidden" is not a good status code for us. It indicates that
something is likely wrong with the access token / authorization. In that
case we want user to refresh the access token and reconnect. However if
we're under heavy load, we want to say user to try again much later.
E.g. maybe return 429 or 503 status code.

This commit allows now returning `{:error, status}` in the
Socket#connect function. It is backwards compatible: just `:error` is
still mapped to Forbidden.

lib/phoenix/socket.ex

@@ -203,8 +203,14 @@ defmodule Phoenix.Socket do
   See `Phoenix.Token` documentation for examples in
   performing token verification on connect.
   """
-  @callback connect(params :: map, Socket.t) :: {:ok, Socket.t} | :error
-  @callback connect(params :: map, Socket.t, connect_info :: map) :: {:ok, Socket.t} | :error
+  @callback connect(params :: map, Socket.t) ::
+              {:ok, Socket.t}
+              | :error
+              | {:error, atom() | integer()}
+  @callback connect(params :: map, Socket.t, connect_info :: map) ::
+              {:ok, Socket.t}
+              | :error
+              | {:error, atom() | integer()}
 
   @doc ~S"""
   Identifies the socket connection.
@@ -456,6 +462,7 @@ defmodule Phoenix.Socket do
 
   defp result({:ok, _}), do: :ok
   defp result(:error), do: :error
+  defp result({:error, _status}), do: :error
 
   def __init__({state, %{id: id, endpoint: endpoint} = socket}) do
     _ = id && endpoint.subscribe(id, link: true)
@@ -589,6 +596,9 @@ defmodule Phoenix.Socket do
             :error
         end
 
+      {:error, status} ->
+        {:error, status}
+
       :error ->
         :error
 

lib/phoenix/transports/long_poll.ex

@@ -100,15 +100,15 @@ defmodule Phoenix.Transports.LongPoll do
       "phx:lp:"
       <> Base.encode64(:crypto.strong_rand_bytes(16))
       <> (System.system_time(:millisecond) |> Integer.to_string)
-      
+
     keys = Keyword.get(opts, :connect_info, [])
     connect_info = Transport.connect_info(conn, endpoint, keys)
     arg = {endpoint, handler, opts, conn.params, priv_topic, connect_info}
     spec = {Phoenix.Transports.LongPoll.Server, arg}
 
     case DynamicSupervisor.start_child(Phoenix.Transports.LongPoll.Supervisor, spec) do
-      :ignore ->
-        conn |> put_status(:forbidden) |> status_json()
+      {:error, status} ->
+        conn |> put_status(status) |> status_json()
 
       {:ok, server_pid} ->
         data  = {:v1, endpoint.config(:endpoint_id), server_pid, priv_topic}

lib/phoenix/transports/long_poll_server.ex

@@ -36,8 +36,11 @@ defmodule Phoenix.Transports.LongPoll.Server do
         schedule_inactive_shutdown(state.window_ms)
         {:ok, state}
 
-    :error ->
-      :ignore
+      {:error, status} ->
+        {:stop, status}
+
+      :error ->
+        {:stop, :forbidden}
     end
   end
 

lib/phoenix/transports/websocket.ex

@@ -31,6 +31,7 @@ defmodule Phoenix.Transports.WebSocket do
 
         case handler.connect(config) do
           {:ok, state} -> {:ok, conn, state}
+          {:error, status} -> {:error, Plug.Conn.send_resp(conn, status, "")}
           :error -> {:error, Plug.Conn.send_resp(conn, 403, "")}
         end
     end

test/phoenix/integration/long_poll_socket_test.exs

@@ -32,7 +32,12 @@ defmodule Phoenix.Integration.LongPollSocketTest do
 
     def connect(map) do
       %{endpoint: Endpoint, params: params, transport: :longpoll} = map
-      {:ok, {:params, params}}
+
+      if (status = Map.get(params, "error")) do
+        {:error, String.to_atom(status)}
+      else
+        {:ok, {:params, params}}
+      end
     end
 
     def init({:params, _} = state) do
@@ -145,4 +150,12 @@ defmodule Phoenix.Integration.LongPollSocketTest do
     resp = poll(:get, "ws/longpoll", secret, nil)
     assert resp.body["messages"] == ["pong"]
   end
+
+  test "returns custom error code" do
+    resp = poll(:get, "ws/longpoll", %{"error" => "service_unavailable"}, nil)
+    assert resp.body["status"] == 503
+
+    resp = poll(:get, "ws/longpoll", %{"error" => "too_many_requests"}, nil)
+    assert resp.body["status"] == 429
+  end
 end

test/phoenix/integration/websocket_socket_test.exs

@@ -30,7 +30,12 @@ defmodule Phoenix.Integration.WebSocketTest do
 
     def connect(map) do
       %{endpoint: Endpoint, params: params, transport: :websocket} = map
-      {:ok, {:params, params}}
+
+      if (status = Map.get(params, "error")) do
+        {:error, String.to_atom(status)}
+      else
+        {:ok, {:params, params}}
+      end
     end
 
     def init({:params, _} = state) do
@@ -113,6 +118,11 @@ defmodule Phoenix.Integration.WebSocketTest do
     assert_receive {:text, "pong"}
   end
 
+  test "returns custom error code" do
+    assert {:error, {503, _}} = WebsocketClient.start_link(self(), "#{@path}?error=service_unavailable", :noop)
+    assert {:error, {429, _}} = WebsocketClient.start_link(self(), "#{@path}?error=too_many_requests", :noop)
+  end
+
   test "allows a custom path" do
     path = "ws://127.0.0.1:#{@port}/custom/some_path/nested/path"
     assert {:ok, _} = WebsocketClient.start_link(self(), "#{path}?key=value", :noop)