Configurable csrf token session key (phoenixframework#5183)

lib/phoenix/endpoint.ex

@@ -817,6 +817,10 @@ defmodule Phoenix.Endpoint do
           config in runtime `{MyAppWeb.Auth, :get_session_config, []}`. Otherwise
           the session will be `nil`.
 
+          `session_config` may take a `:csrf_token_key` option
+          which is useful when using `:protect_from_forgery` with a custom
+          `:session_key`. If not given, it defaults to `"_csrf_token"`.
+
       Arbitrary keywords may also appear following the above valid keys, which
       is useful for passing custom connection information to the socket.
 

lib/phoenix/socket/transport.ex

@@ -259,11 +259,14 @@ defmodule Phoenix.Socket.Transport do
     [connect_info: connect_info] ++ config
   end
 
+  # The original session_config is returned in addition to init value so we can
+  # access special config like :csrf_token_key downstream.
   defp init_session(session_config) when is_list(session_config) do
     key = Keyword.fetch!(session_config, :key)
     store = Plug.Session.Store.get(Keyword.fetch!(session_config, :store))
     init = store.init(Keyword.drop(session_config, [:store, :key]))
-    {key, store, init}
+    csrf_token_key = Keyword.get(session_config, :csrf_token_key, "_csrf_token")
+    {key, store, {csrf_token_key, init}}
   end
 
   defp init_session({_, _, _} = mfa) do
@@ -465,15 +468,15 @@ defmodule Phoenix.Socket.Transport do
     end
   end
 
-  defp connect_session(conn, endpoint, {key, store, store_config}) do
+  defp connect_session(conn, endpoint, {key, store, {csrf_token_key, init}}) do
     conn = Plug.Conn.fetch_cookies(conn)
 
     with csrf_token when is_binary(csrf_token) <- conn.params["_csrf_token"],
          cookie when is_binary(cookie) <- conn.cookies[key],
          conn = put_in(conn.secret_key_base, endpoint.config(:secret_key_base)),
-         {_, session} <- store.get(conn, cookie, store_config),
+         {_, session} <- store.get(conn, cookie, init),
          csrf_state when is_binary(csrf_state) <-
-           Plug.CSRFProtection.dump_state_from_session(session["_csrf_token"]),
+          Plug.CSRFProtection.dump_state_from_session(session[csrf_token_key]),
          true <- Plug.CSRFProtection.valid_state_and_csrf_token?(csrf_state, csrf_token) do
       session
     else

test/phoenix/socket/transport_test.exs

@@ -23,13 +23,20 @@ defmodule Phoenix.Socket.TransportTest do
       signing_salt: "change_me"
     ]
 
-    def session_config, do: @session_config
+    def session_config(overrides \\ []), do: Keyword.merge(@session_config, overrides)
 
     plug Plug.Session, @session_config
     plug :fetch_session
-    plug Plug.CSRFProtection
+    plug :put_csrf
     plug :put_session
 
+    defp put_csrf(conn, _opts) do
+      conn = Plug.Conn.fetch_query_params(conn)
+      session_key = Map.get(conn.query_params, "session_key", "_csrf_token")
+      opts = Plug.CSRFProtection.init(session_key: session_key)
+      Plug.CSRFProtection.call(conn, opts)
+    end
+
     defp put_session(conn, _) do
       conn
       |> put_session(:from_session, "123")
@@ -288,5 +295,40 @@ defmodule Phoenix.Socket.TransportTest do
                |> fetch_query_params()
                |> Transport.connect_info(Endpoint, connect_info)
     end
+
+    test "loads the session with custom :csrf_token_key" do
+      conn = conn(:get, "https://foo.com?session_key=_custom_csrf_token") |> Endpoint.call([])
+      csrf_token = conn.resp_body
+      session_cookie = conn.cookies["_hello_key"]
+
+      connect_info = load_connect_info(
+        session: {
+          Endpoint,
+          :session_config,
+          [[csrf_token_key: "_custom_csrf_token"]]
+        }
+      )
+
+      assert %{session: %{"from_session" => "123"}} =
+              conn(:get, "https://foo.com/", _csrf_token: csrf_token)
+              |> put_req_cookie("_hello_key", session_cookie)
+              |> fetch_query_params()
+              |> Transport.connect_info(Endpoint, connect_info)
+
+      connect_info = load_connect_info(
+        session: {
+          Endpoint,
+          :session_config,
+          [[csrf_token_key: "bad_key"]]
+        }
+      )
+
+      assert %{session: nil} =
+              conn(:get, "https://foo.com/", _csrf_token: csrf_token)
+              |> put_req_cookie("_hello_key", session_cookie)
+              |> fetch_query_params()
+              |> Transport.connect_info(Endpoint, connect_info)
+    end
+
   end
 end