support refreshing connection config during reconnect

examples/token_refresh/README.md

@@ -2,14 +2,12 @@
 
 Often times, the socket server you want to connect to requires some form
 of token authentication. Sometimes these tokens expire and need to be
-refreshed before a reconnect happens. `Slipstream.connect/2` is helpful
-to use in place of `Slipstream.reconnect/1` if you need to modify your
-connection configuration before reconnecting.
+refreshed before a reconnect happens. The `Slipstream.refresh_connection_config/2` callback is there to make it easy to update connection configuration before reconnecting.
 
 ## Tutorial
 
 In this tutorial, we'll build a basic Slipstream client that can handle
-a 403 from our websocket and reconnect successfully.
+refresh a disconnection from our websocket and reconnect successfully.
 
 Let's begin with a fresh client
 
@@ -34,71 +32,42 @@ end
 ```
 
 Notice that we simply `reconnect/2` inside of the `handle_disconnect/2`
-callback. This is the perfect place to handle an expired token. The 
-`reason` value will contain information about why we disconnected. In
-this case, we're looking for a Mint error with a `403` status code. The 
-error will roughly look like this:
+callback.
 
-```elixir
-{:error, {_upgrade_failure, %{status_code: 403}}}
-```
-
-Let's update our `handle_disconnect/2` callback to handle that situation.
+Now that we know when a token needs to be refreshed, let's add in the callback to our client.
 
 ```elixir
 @impl Slipstream
-def handle_disconnect({:error, {_, %{status_code: 403}}}, socket) do
-  # get new token and then attempt to reconnect
+def refresh_connection_config(_socket, config) do
+  update_uri(config)
 end
-```
-
-Now that we know when a token needs to be refreshed, let's add in some
-basic token logic to our client. First, we'll wrap `connect/2` with a
-custom function that retrieves a token and modifies our config:
-
-```elixir
-defp make_new_token, do: "your_token_logic"
-
-defp connect_with_token(socket) do
-  new_token = make_new_token()
-
-  socket =
-    update(socket, :config, fn config ->
-      uri =
-        config
-        |> Keyword.get(:uri)
-        |> URI.parse()
-        |> Map.put(:query, "token=#{new_token}")
-        |> URI.to_string()
 
-      Keyword.put(config, :uri, uri)
-    end)
+defp update_uri(config) do
+  uri =
+    config
+    |> Keyword.get(:uri)
+    |> URI.parse()
+    |> Map.put(:query, "token=#{make_new_token()}")
+    |> URI.to_string()
 
-  connect(socket, socket.assigns.config)
+  Keyword.put(config, :uri, uri)
 end
+
+defp make_new_token(), do: "get_new_token_here"
 ```
 
 The `config` on the socket can be unwrapped with the `URI` module. Then, we
 simply append our token onto the query and change it back into a string.
 This gives us a clean API to ensure our connections always contain a token.
-Let's update our `init/1` and `handle_disconnect/2` callbacks to use this
-new function:
+Let's update our `init/1` to use this new function:
 
 ```elixir
 def init(config) do
+  config = update_uri(config)
+
   new_socket()
   |> assign(:config, config)
-  |> connect_with_token()
-end
-
-@impl Slipstream
-def handle_disconnect({:error, {_, %{status_code: 403}}}, socket) do
-  connect_with_token(socket)
-end
-
-@impl Slipstream
-def handle_disconnect(_reason, socket) do
-  reconnect(socket)
+  |> connect(config)
 end
 ```
 

examples/token_refresh/client.ex

@@ -11,38 +11,33 @@ defmodule MyApp.TokenRefreshClient do
 
   @impl Slipstream
   def init(config) do
+    config = update_uri(config)
+
     new_socket()
     |> assign(:config, config)
-    |> connect_with_token()
+    |> connect(config)
   end
 
-  defp make_new_token, do: "get_new_token_here"
-
-  defp connect_with_token(socket) do
-    new_token = make_new_token()
-
-    socket =
-      update(socket, :config, fn config ->
-        uri =
-          config
-          |> Keyword.get(:uri)
-          |> URI.parse()
-          |> Map.put(:query, "token=#{new_token}")
-          |> URI.to_string()
-
-        Keyword.put(config, :uri, uri)
-      end)
-
-    connect(socket, socket.assigns.config)
+  @impl Slipstream
+  def handle_disconnect(_reason, socket) do
+    reconnect(socket)
   end
 
   @impl Slipstream
-  def handle_disconnect({:error, {_, %{status_code: 403}}}, socket) do
-    connect_with_token(socket)
+  def refresh_connection_config(_socket, config) do
+    update_uri(config)
   end
 
-  @impl Slipstream
-  def handle_disconnect(_reason, socket) do
-    reconnect(socket)
+  defp update_uri(config) do
+    uri =
+      config
+      |> Keyword.get(:uri)
+      |> URI.parse()
+      |> Map.put(:query, "token=#{make_new_token()}")
+      |> URI.to_string()
+
+    Keyword.put(config, :uri, uri)
   end
+
+  defp make_new_token(), do: "get_new_token_here"
 end

lib/slipstream.ex

@@ -226,7 +226,7 @@ defmodule Slipstream do
       end
   """
 
-  alias Slipstream.{Commands, Events, Socket, TelemetryHelper}
+  alias Slipstream.{Commands, Events, Socket, TelemetryHelper, Configuration}
   import Slipstream.CommandRouter, only: [route_command: 1]
   import Slipstream.Signatures, only: [event: 1, command: 1]
 
@@ -694,6 +694,28 @@ defmodule Slipstream do
               | {:stop, stop_reason :: term(), new_socket}
             when new_socket: Socket.t()
 
+  @doc """
+  Invoked when a socket reconnection is triggered via `reconnect/1`.
+
+  This callback allows for connection config properties to be updated,
+  useful for cases where authentication tokens or headers might have
+  expired.
+
+  ## Examples
+
+      @impl Slipstream
+      def refresh_connection_config(socket, config) do
+        config
+        |> Map.put(:headers, generate_headers(socket))
+        |> Map.put(:url, update_uri(socket))
+      end
+  """
+  @doc since: "1.1.2"
+  @callback refresh_connection_config(
+              socket :: Socket.t(),
+              config :: Configuration.t()
+            ) :: Configuration.t()
+
   @optional_callbacks init: 1,
                       handle_info: 2,
                       handle_cast: 2,
@@ -706,7 +728,8 @@ defmodule Slipstream do
                       handle_message: 4,
                       handle_reply: 3,
                       handle_topic_close: 3,
-                      handle_leave: 2
+                      handle_leave: 2,
+                      refresh_connection_config: 2
 
   # --- core functionality
 
@@ -1508,8 +1531,6 @@ defmodule Slipstream do
         |> Supervisor.child_spec(unquote(Macro.escape(opts)))
       end
 
-      defoverridable child_spec: 1
-
       require Slipstream.Signatures
 
       import Slipstream
@@ -1532,6 +1553,9 @@ defmodule Slipstream do
         Slipstream.Callback.dispatch(__MODULE__, event, socket)
       end
 
+      @impl Slipstream
+      def refresh_connection_config(socket, config), do: config
+
       # this matches on time-delay commands like those emitted from
       # reconnect/1 and rejoin/3
       def handle_info(
@@ -1540,7 +1564,11 @@ defmodule Slipstream do
             ),
             socket
           ) do
-        socket = TelemetryHelper.begin_connect(socket, cmd.config)
+        config = refresh_connection_config(socket, cmd.config)
+
+        socket = TelemetryHelper.begin_connect(socket, config)
+
+        cmd = %{cmd | config: config}
 
         _ = Slipstream.CommandRouter.route_command(cmd)
 
@@ -1559,6 +1587,8 @@ defmodule Slipstream do
 
         {:noreply, socket}
       end
+
+      defoverridable child_spec: 1, refresh_connection_config: 2
     end
   end
 end