Fetch session from conn

Author: LostKobrakai

lib/phoenix/endpoint.ex

@@ -643,7 +643,7 @@ defmodule Phoenix.Endpoint do
   end
 
   @doc false
-  defmacro __before_compile__(%{module: module}) do
+  defmacro __before_compile__(_env) do
     quote do
       defoverridable call: 2
 

lib/phoenix/endpoint/supervisor.ex

@@ -84,9 +84,9 @@ defmodule Phoenix.Endpoint.Supervisor do
       config_children(mod, secret_conf, default_conf) ++
         warmup_children(mod) ++
         pubsub_children(mod, conf) ++
-        socket_children(mod, conf, :child_spec) ++
+        socket_children(mod, :child_spec) ++
         server_children(mod, conf, server?) ++
-        socket_children(mod, conf, :drainer_spec) ++
+        socket_children(mod, :drainer_spec) ++
         watcher_children(mod, conf, server?)
 
     Supervisor.init(children, strategy: :one_for_one)
@@ -118,7 +118,7 @@ defmodule Phoenix.Endpoint.Supervisor do
     end
   end
 
-  defp socket_children(endpoint, conf, fun) do
+  defp socket_children(endpoint, fun) do
     for {socket, opts} <- endpoint.__start_sockets__(),
         # TODO is this the correct place for this?
         # Needs to know transport specific config

lib/phoenix/socket/transport.ex

@@ -290,39 +290,23 @@ defmodule Phoenix.Socket.Transport do
                :user_agent,
                :x_headers,
                :sec_websocket_headers,
-               :auth_token
+               :auth_token,
+               :session
              ] ->
           key
 
-        {:session, session} ->
-          {:session, init_session(session)}
-
         {_, _} = pair ->
           pair
 
         other ->
           raise ArgumentError,
-                ":connect_info keys are expected to be one of :peer_data, :trace_context_headers, :x_headers, :user_agent, :sec_websocket_headers, :uri, or {:session, config}, " <>
+                ":connect_info keys are expected to be one of :peer_data, :trace_context_headers, :x_headers, :user_agent, :sec_websocket_headers, :uri, or :session, " <>
                   "optionally followed by custom keyword pairs, got: #{inspect(other)}"
       end)
 
     [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]))
-    csrf_token_key = Keyword.get(session_config, :csrf_token_key, "_csrf_token")
-    {key, store, {csrf_token_key, init}}
-  end
-
-  defp init_session({_, _, _} = mfa) do
-    {:mfa, mfa}
-  end
-
   @doc """
   Runs the code reloader if enabled.
   """
@@ -496,7 +480,7 @@ defmodule Phoenix.Socket.Transport do
 
   The CSRF check can be disabled by setting the `:check_csrf` option to `false`.
   """
-  def connect_info(conn, endpoint, keys, opts \\ []) do
+  def connect_info(conn, _endpoint, keys, opts \\ []) do
     for key <- keys, into: %{} do
       case key do
         :peer_data ->
@@ -517,8 +501,8 @@ defmodule Phoenix.Socket.Transport do
         :sec_websocket_headers ->
           {:sec_websocket_headers, fetch_headers(conn, "sec-websocket-")}
 
-        {:session, session} ->
-          {:session, connect_session(conn, endpoint, session, opts)}
+        :session ->
+          {:session, connect_session(conn, opts)}
 
         :auth_token ->
           {:auth_token, conn.private[:phoenix_transport_auth_token]}
@@ -529,28 +513,15 @@ defmodule Phoenix.Socket.Transport do
     end
   end
 
-  defp connect_session(conn, endpoint, {key, store, {csrf_token_key, init}}, opts) do
-    conn = Plug.Conn.fetch_cookies(conn)
+  defp connect_session(conn, opts) do
+    session = Plug.Conn.get_session(conn)
     check_csrf = Keyword.get(opts, :check_csrf, true)
+    csrf_token_key = Keyword.get(opts, :csrf_token_key, "_csrf_token")
 
-    with 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, init),
-         true <- not check_csrf or csrf_token_valid?(conn, session, csrf_token_key) do
+    if not check_csrf or csrf_token_valid?(conn, session, csrf_token_key) do
       session
     else
-      _ -> nil
-    end
-  end
-
-  defp connect_session(conn, endpoint, {:mfa, {module, function, args}}, opts) do
-    case apply(module, function, args) do
-      session_config when is_list(session_config) ->
-        connect_session(conn, endpoint, init_session(session_config), opts)
-
-      other ->
-        raise ArgumentError,
-              "the MFA given to `session_config` must return a keyword list, got: #{inspect(other)}"
+      nil
     end
   end
 

test/phoenix/integration/websocket_channels_test.exs

@@ -162,52 +162,72 @@ defmodule Phoenix.Integration.WebSocketChannelsTest do
     def handle_error(conn, :rate_limit), do: Plug.Conn.send_resp(conn, 429, "Too many requests")
   end
 
-  defmodule Endpoint do
-    use Phoenix.Endpoint, otp_app: :phoenix
+  defmodule SetSession do
+    import Plug.Conn
 
-    @session_config store: :cookie,
-                    key: "_hello_key",
-                    signing_salt: "change_me"
+    def init(opts), do: opts
 
-    socket "/ws", UserSocket,
-      websocket: [
-        check_origin: ["//example.com"],
-        timeout: 200,
-        error_handler: {UserSocket, :handle_error, []}
-      ]
+    def call(conn, _) do
+      conn
+      |> put_session(:from_session, "123")
+      |> send_resp(200, Plug.CSRFProtection.get_csrf_token())
+    end
+  end
 
-    socket "/ws/admin", UserSocket,
-      websocket: [
-        check_origin: ["//example.com"],
-        timeout: 200
-      ]
+  defmodule Router do
+    use Phoenix.Router
+    import Phoenix.Socket.Router
+
+    get "/", SetSession, []
+
+    scope "/ws" do
+      socket "/", UserSocket,
+        websocket: [
+          check_origin: ["//example.com"],
+          timeout: 200,
+          error_handler: {UserSocket, :handle_error, []}
+        ]
+
+      socket "/admin", UserSocket,
+        websocket: [
+          check_origin: ["//example.com"],
+          timeout: 200
+        ]
 
-    socket "/ws/connect_info", UserSocketConnectInfo,
-      websocket: [
-        check_origin: ["//example.com"],
-        timeout: 200,
-        connect_info: [
-          :trace_context_headers,
-          :x_headers,
-          :peer_data,
-          :uri,
-          :user_agent,
-          :sec_websocket_headers,
-          session: @session_config,
-          signing_salt: "salt",
+      socket "/connect_info", UserSocketConnectInfo,
+        websocket: [
+          check_origin: ["//example.com"],
+          timeout: 200,
+          connect_info: [
+            :trace_context_headers,
+            :x_headers,
+            :peer_data,
+            :uri,
+            :user_agent,
+            :sec_websocket_headers,
+            :session,
+            signing_salt: "salt"
+          ]
         ]
+    end
+  end
+
+  defmodule Endpoint do
+    use Phoenix.Endpoint,
+      otp_app: :phoenix,
+      start_sockets: [
+        {UserSocket, []},
+        {UserSocketConnectInfo, []}
       ]
 
-    plug Plug.Session, @session_config
+    plug Plug.Session,
+      store: :cookie,
+      key: "_hello_key",
+      signing_salt: "change_me"
+
     plug :fetch_session
     plug Plug.CSRFProtection
-    plug :put_session
-
-    defp put_session(conn, _) do
-      conn
-      |> put_session(:from_session, "123")
-      |> send_resp(200, Plug.CSRFProtection.get_csrf_token())
-    end
+    plug Router
   end
 
   setup %{adapter: adapter} do
@@ -540,7 +560,7 @@ defmodule Phoenix.Integration.WebSocketChannelsTest do
               {:ok, _} = WebsocketClient.connect(self(), @vsn_path, @serializer)
             end)
 
-          assert log == ""
+          refute log =~ "UserSocket"
         end
 
         test "logs and filter params on join and handle_in" do