feat(realtime): pull access token mechanism (#615)

* feat(realtime): pull access token mechanism

* add tests

* report issue if custom access token is assigned

* add docs

* pull token from third party or fallback to auth

Author: grdsdev

Sources/Realtime/V2/RealtimeChannelV2.swift

@@ -29,7 +29,7 @@ struct Socket: Sendable {
   var broadcastURL: @Sendable () -> URL
   var status: @Sendable () -> RealtimeClientStatus
   var options: @Sendable () -> RealtimeClientOptions
-  var accessToken: @Sendable () -> String?
+  var accessToken: @Sendable () async -> String?
   var apiKey: @Sendable () -> String?
   var makeRef: @Sendable () -> Int
 
@@ -46,7 +46,12 @@ extension Socket {
       broadcastURL: { [weak client] in client?.broadcastURL ?? URL(string: "http://localhost")! },
       status: { [weak client] in client?.status ?? .disconnected },
       options: { [weak client] in client?.options ?? .init() },
-      accessToken: { [weak client] in client?.mutableState.accessToken },
+      accessToken: { [weak client] in
+        if let accessToken = try? await client?.options.accessToken?() {
+          return accessToken
+        }
+        return client?.mutableState.accessToken
+      },
       apiKey: { [weak client] in client?.apikey },
       makeRef: { [weak client] in client?.makeRef() ?? 0 },
       connect: { [weak client] in await client?.connect() },
@@ -139,7 +144,7 @@ public final class RealtimeChannelV2: Sendable {
 
     let payload = RealtimeJoinPayload(
       config: joinConfig,
-      accessToken: socket.accessToken()
+      accessToken: await socket.accessToken()
     )
 
     let joinRef = socket.makeRef().description
@@ -213,7 +218,7 @@ public final class RealtimeChannelV2: Sendable {
       if let apiKey = socket.apiKey() {
         headers[.apiKey] = apiKey
       }
-      if let accessToken = socket.accessToken() {
+      if let accessToken = await socket.accessToken() {
         headers[.authorization] = "Bearer \(accessToken)"
       }
 

Sources/Realtime/V2/RealtimeClientV2.swift

@@ -107,7 +107,11 @@ public final class RealtimeClientV2: Sendable {
     apikey = options.apikey
 
     mutableState.withValue {
-      $0.accessToken = options.accessToken ?? options.apikey
+      if let accessToken = options.headers[.authorization]?.split(separator: " ").last {
+        $0.accessToken = String(accessToken)
+      } else {
+        $0.accessToken = options.apikey
+      }
     }
   }
 
@@ -361,8 +365,22 @@ public final class RealtimeClientV2: Sendable {
   }
 
   /// Sets the JWT access token used for channel subscription authorization and Realtime RLS.
-  /// - Parameter token: A JWT string.
-  public func setAuth(_ token: String?) async {
+  ///
+  /// If `token` is nil it will use the ``RealtimeClientOptions/accessToken`` callback function or the token set on the client.
+  ///
+  /// On callback used, it will set the value of the token internal to the client.
+  /// - Parameter token: A JWT string to override the token set on the client.
+  public func setAuth(_ token: String? = nil) async {
+    var token = token
+
+    if token == nil {
+      token = try? await options.accessToken?()
+    }
+
+    if token == nil {
+      token = mutableState.accessToken
+    }
+
     if let token, let payload = JWT.decodePayload(token),
       let exp = payload["exp"] as? TimeInterval, exp < Date().timeIntervalSince1970
     {
@@ -371,7 +389,7 @@ public final class RealtimeClientV2: Sendable {
       return
     }
 
-    mutableState.withValue {
+    mutableState.withValue { [token] in
       $0.accessToken = token
     }
 

Sources/Realtime/V2/Types.swift

@@ -6,8 +6,8 @@
 //
 
 import Foundation
-import Helpers
 import HTTPTypes
+import Helpers
 
 #if canImport(FoundationNetworking)
   import FoundationNetworking
@@ -22,6 +22,7 @@ public struct RealtimeClientOptions: Sendable {
   var disconnectOnSessionLoss: Bool
   var connectOnSubscribe: Bool
   var fetch: (@Sendable (_ request: URLRequest) async throws -> (Data, URLResponse))?
+  package var accessToken: (@Sendable () async throws -> String)?
   package var logger: (any SupabaseLogger)?
 
   public static let defaultHeartbeatInterval: TimeInterval = 15
@@ -38,6 +39,7 @@ public struct RealtimeClientOptions: Sendable {
     disconnectOnSessionLoss: Bool = Self.defaultDisconnectOnSessionLoss,
     connectOnSubscribe: Bool = Self.defaultConnectOnSubscribe,
     fetch: (@Sendable (_ request: URLRequest) async throws -> (Data, URLResponse))? = nil,
+    accessToken: (@Sendable () async throws -> String)? = nil,
     logger: (any SupabaseLogger)? = nil
   ) {
     self.headers = HTTPFields(headers)
@@ -47,19 +49,13 @@ public struct RealtimeClientOptions: Sendable {
     self.disconnectOnSessionLoss = disconnectOnSessionLoss
     self.connectOnSubscribe = connectOnSubscribe
     self.fetch = fetch
+    self.accessToken = accessToken
     self.logger = logger
   }
 
   var apikey: String? {
     headers[.apiKey]
   }
-
-  var accessToken: String? {
-    guard let accessToken = headers[.authorization]?.split(separator: " ").last else {
-      return nil
-    }
-    return String(accessToken)
-  }
 }
 
 public typealias RealtimeSubscription = ObservationToken

Sources/Supabase/SupabaseClient.swift

@@ -2,12 +2,12 @@
 import ConcurrencyExtras
 import Foundation
 @_exported import Functions
+import HTTPTypes
 import Helpers
 import IssueReporting
 @_exported import PostgREST
 @_exported import Realtime
 @_exported import Storage
-import HTTPTypes
 
 #if canImport(FoundationNetworking)
   import FoundationNetworking
@@ -33,10 +33,11 @@ public final class SupabaseClient: Sendable {
   /// Supabase Auth allows you to create and manage user sessions for access to data that is secured by access policies.
   public var auth: AuthClient {
     if options.auth.accessToken != nil {
-      reportIssue("""
-      Supabase Client is configured with the auth.accessToken option,
-      accessing supabase.auth is not possible.
-      """)
+      reportIssue(
+        """
+        Supabase Client is configured with the auth.accessToken option,
+        accessing supabase.auth is not possible.
+        """)
     }
     return _auth
   }
@@ -80,7 +81,14 @@ public final class SupabaseClient: Sendable {
   let _realtime: UncheckedSendable<RealtimeClient>
 
   /// Realtime client for Supabase
-  public let realtimeV2: RealtimeClientV2
+  public var realtimeV2: RealtimeClientV2 {
+    mutableState.withValue {
+      if $0.realtime == nil {
+        $0.realtime = _initRealtimeClient()
+      }
+      return $0.realtime!
+    }
+  }
 
   /// Supabase Functions allows you to deploy and invoke edge functions.
   public var functions: FunctionsClient {
@@ -112,6 +120,7 @@ public final class SupabaseClient: Sendable {
     var storage: SupabaseStorageClient?
     var rest: PostgrestClient?
     var functions: FunctionsClient?
+    var realtime: RealtimeClientV2?
 
     var changedAccessToken: String?
   }
@@ -189,18 +198,6 @@ public final class SupabaseClient: Sendable {
       )
     )
 
-    var realtimeOptions = options.realtime
-    realtimeOptions.headers.merge(with: _headers)
-
-    if realtimeOptions.logger == nil {
-      realtimeOptions.logger = options.global.logger
-    }
-
-    realtimeV2 = RealtimeClientV2(
-      url: supabaseURL.appendingPathComponent("/realtime/v1"),
-      options: realtimeOptions
-    )
-
     if options.auth.accessToken == nil {
       listenForAuthEvents()
     }
@@ -351,11 +348,7 @@ public final class SupabaseClient: Sendable {
   }
 
   private func adapt(request: URLRequest) async -> URLRequest {
-    let token: String? = if let accessToken = options.auth.accessToken {
-      try? await accessToken()
-    } else {
-      try? await auth.session.accessToken
-    }
+    let token = try? await _getAccessToken()
 
     var request = request
     if let token {
@@ -364,6 +357,14 @@ public final class SupabaseClient: Sendable {
     return request
   }
 
+  private func _getAccessToken() async throws -> String {
+    if let accessToken = options.auth.accessToken {
+      try await accessToken()
+    } else {
+      try await auth.session.accessToken
+    }
+  }
+
   private func listenForAuthEvents() {
     let task = Task {
       for await (event, session) in auth.authStateChanges {
@@ -377,7 +378,9 @@ public final class SupabaseClient: Sendable {
 
   private func handleTokenChanged(event: AuthChangeEvent, session: Session?) async {
     let accessToken: String? = mutableState.withValue {
-      if [.initialSession, .signedIn, .tokenRefreshed].contains(event), $0.changedAccessToken != session?.accessToken {
+      if [.initialSession, .signedIn, .tokenRefreshed].contains(event),
+        $0.changedAccessToken != session?.accessToken
+      {
         $0.changedAccessToken = session?.accessToken
         return session?.accessToken ?? supabaseKey
       }
@@ -393,4 +396,33 @@ public final class SupabaseClient: Sendable {
     realtime.setAuth(accessToken)
     await realtimeV2.setAuth(accessToken)
   }
+
+  private func _initRealtimeClient() -> RealtimeClientV2 {
+    var realtimeOptions = options.realtime
+    realtimeOptions.headers.merge(with: _headers)
+
+    if realtimeOptions.logger == nil {
+      realtimeOptions.logger = options.global.logger
+    }
+
+    if realtimeOptions.accessToken == nil {
+      realtimeOptions.accessToken = { [weak self] in
+        try await self?._getAccessToken() ?? ""
+      }
+    } else {
+      reportIssue(
+        """
+        You assigned a custom `accessToken` closure to the RealtimeClientV2. This might not work as you expect
+        as SupabaseClient uses Auth for pulling an access token to send on the realtime channels.
+
+        Please make sure you know what you're doing.
+        """
+      )
+    }
+
+    return RealtimeClientV2(
+      url: supabaseURL.appendingPathComponent("/realtime/v1"),
+      options: realtimeOptions
+    )
+  }
 }

Tests/RealtimeTests/RealtimeTests.swift

@@ -36,7 +36,10 @@ final class RealtimeTests: XCTestCase {
         headers: ["apikey": apiKey],
         heartbeatInterval: 1,
         reconnectDelay: 1,
-        timeoutInterval: 2
+        timeoutInterval: 2,
+        accessToken: {
+          "custom.access.token"
+        }
       ),
       ws: ws,
       http: http
@@ -100,7 +103,7 @@ final class RealtimeTests: XCTestCase {
           "event" : "phx_join",
           "join_ref" : "1",
           "payload" : {
-            "access_token" : "anon.api.key",
+            "access_token" : "custom.access.token",
             "config" : {
               "broadcast" : {
                 "ack" : false,
@@ -179,7 +182,7 @@ final class RealtimeTests: XCTestCase {
           "event" : "phx_join",
           "join_ref" : "1",
           "payload" : {
-            "access_token" : "anon.api.key",
+            "access_token" : "custom.access.token",
             "config" : {
               "broadcast" : {
                 "ack" : false,
@@ -201,7 +204,7 @@ final class RealtimeTests: XCTestCase {
           "event" : "phx_join",
           "join_ref" : "2",
           "payload" : {
-            "access_token" : "anon.api.key",
+            "access_token" : "custom.access.token",
             "config" : {
               "broadcast" : {
                 "ack" : false,
@@ -322,7 +325,7 @@ final class RealtimeTests: XCTestCase {
     assertInlineSnapshot(of: request?.urlRequest, as: .raw(pretty: true)) {
       """
       POST https://localhost:54321/realtime/v1/api/broadcast
-      Authorization: Bearer anon.api.key
+      Authorization: Bearer custom.access.token
       Content-Type: application/json
       apiKey: anon.api.key