[PM-23189] Add ClientManagedTokens trait (#337)

Adds support for client managed access tokens alongside SDK managed
access tokens.

It does so by splitting `Tokens` into `SdkManaged` and `ClientManaged`.
`ClientManaged` relies on the consumer implementing a trait with
`get_access_token` function.

In WASM this is exposed as a `TokenProvider` interface which can be
implemented in TS clients using:

```ts
class JsTokenProvider implements TokenProvider {
  constructor(private apiService: ApiService) {}

  async get_access_token(): Promise<string> {
    return await this.apiService.getActiveBearerToken();
  }
}
```

Author: Hinton

Cargo.lock

@@ -29,6 +29,7 @@ wasm = [
 ] # WASM support
 
 [dependencies]
+async-trait = { workspace = true }
 base64 = ">=0.22.1, <0.23"
 bitwarden-api-api = { workspace = true }
 bitwarden-api-identity = { workspace = true }

crates/bitwarden-core/Cargo.toml

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use chrono::Utc;
 
 use super::login::LoginError;
@@ -10,18 +12,44 @@ use crate::{
 };
 use crate::{
     auth::api::{request::ApiTokenRequest, response::IdentityTokenResponse},
-    client::{internal::InternalClient, LoginMethod, UserLoginMethod},
+    client::{
+        internal::{ClientManagedTokens, InternalClient, SdkManagedTokens, Tokens},
+        LoginMethod, UserLoginMethod,
+    },
     NotAuthenticatedError,
 };
 
 pub(crate) async fn renew_token(client: &InternalClient) -> Result<(), LoginError> {
-    const TOKEN_RENEW_MARGIN_SECONDS: i64 = 5 * 60;
-
-    let tokens = client
+    let tokens_guard = client
         .tokens
         .read()
         .expect("RwLock is not poisoned")
         .clone();
+
+    match tokens_guard {
+        Tokens::SdkManaged(tokens) => renew_token_sdk_managed(client, tokens).await,
+        Tokens::ClientManaged(tokens) => renew_token_client_managed(client, tokens).await,
+    }
+}
+
+async fn renew_token_client_managed(
+    client: &InternalClient,
+    tokens: Arc<dyn ClientManagedTokens>,
+) -> Result<(), LoginError> {
+    let token = tokens
+        .get_access_token()
+        .await
+        .ok_or(NotAuthenticatedError)?;
+    client.set_api_tokens_internal(token);
+    Ok(())
+}
+
+async fn renew_token_sdk_managed(
+    client: &InternalClient,
+    tokens: SdkManagedTokens,
+) -> Result<(), LoginError> {
+    const TOKEN_RENEW_MARGIN_SECONDS: i64 = 5 * 60;
+
     let login_method = client
         .login_method
         .read()

crates/bitwarden-core/src/auth/renew.rs

@@ -10,7 +10,7 @@ use super::internal::InternalClient;
 use crate::client::flags::Flags;
 use crate::client::{
     client_settings::ClientSettings,
-    internal::{ApiConfigurations, Tokens},
+    internal::{ApiConfigurations, ClientManagedTokens, SdkManagedTokens, Tokens},
 };
 
 /// The main struct to interact with the Bitwarden SDK.
@@ -25,8 +25,20 @@ pub struct Client {
 }
 
 impl Client {
-    #[allow(missing_docs)]
-    pub fn new(settings_input: Option<ClientSettings>) -> Self {
+    /// Create a new Bitwarden client with SDK-managed tokens.
+    pub fn new(settings: Option<ClientSettings>) -> Self {
+        Self::new_internal(settings, Tokens::SdkManaged(SdkManagedTokens::default()))
+    }
+
+    /// Create a new Bitwarden client with client-managed tokens.
+    pub fn new_with_client_tokens(
+        settings: Option<ClientSettings>,
+        tokens: Arc<dyn ClientManagedTokens>,
+    ) -> Self {
+        Self::new_internal(settings, Tokens::ClientManaged(tokens))
+    }
+
+    fn new_internal(settings_input: Option<ClientSettings>, tokens: Tokens) -> Self {
         let settings = settings_input.unwrap_or_default();
 
         fn new_client_builder() -> reqwest::ClientBuilder {
@@ -81,7 +93,7 @@ impl Client {
         Self {
             internal: Arc::new(InternalClient {
                 user_id: OnceLock::new(),
-                tokens: RwLock::new(Tokens::default()),
+                tokens: RwLock::new(tokens),
                 login_method: RwLock::new(None),
                 #[cfg(feature = "internal")]
                 flags: RwLock::new(Flags::default()),

crates/bitwarden-core/src/client/client.rs

@@ -33,11 +33,26 @@ pub struct ApiConfigurations {
     pub device_type: DeviceType,
 }
 
+/// Access and refresh tokens used for authentication and authorization.
+#[derive(Debug, Clone)]
+pub(crate) enum Tokens {
+    SdkManaged(SdkManagedTokens),
+    ClientManaged(Arc<dyn ClientManagedTokens>),
+}
+
+/// Access tokens managed by client applications, such as the web or mobile apps.
+#[async_trait::async_trait]
+pub trait ClientManagedTokens: std::fmt::Debug + Send + Sync {
+    /// Returns the access token, if available.
+    async fn get_access_token(&self) -> Option<String>;
+}
+
+/// Tokens managed by the SDK, the SDK will automatically handle token renewal.
 #[derive(Debug, Default, Clone)]
-pub(crate) struct Tokens {
+pub(crate) struct SdkManagedTokens {
     // These two fields are always written to, but they are not read
     // from the secrets manager SDK.
-    #[cfg_attr(not(feature = "internal"), allow(dead_code))]
+    #[allow(dead_code)]
     access_token: Option<String>,
     pub(crate) expires_on: Option<i64>,
 
@@ -117,11 +132,17 @@ impl InternalClient {
     }
 
     pub(crate) fn set_tokens(&self, token: String, refresh_token: Option<String>, expires_in: u64) {
-        *self.tokens.write().expect("RwLock is not poisoned") = Tokens {
-            access_token: Some(token.clone()),
-            expires_on: Some(Utc::now().timestamp() + expires_in as i64),
-            refresh_token,
-        };
+        *self.tokens.write().expect("RwLock is not poisoned") =
+            Tokens::SdkManaged(SdkManagedTokens {
+                access_token: Some(token.clone()),
+                expires_on: Some(Utc::now().timestamp() + expires_in as i64),
+                refresh_token,
+            });
+        self.set_api_tokens_internal(token);
+    }
+
+    /// Sets api tokens for only internal API clients, use `set_tokens` for SdkManagedTokens.
+    pub(crate) fn set_api_tokens_internal(&self, token: String) {
         let mut guard = self
             .__api_configurations
             .write()
@@ -132,24 +153,6 @@ impl InternalClient {
         inner.api.oauth_access_token = Some(token);
     }
 
-    #[allow(missing_docs)]
-    #[cfg(feature = "internal")]
-    pub fn is_authed(&self) -> bool {
-        let is_token_set = self
-            .tokens
-            .read()
-            .expect("RwLock is not poisoned")
-            .access_token
-            .is_some();
-        let is_login_method_set = self
-            .login_method
-            .read()
-            .expect("RwLock is not poisoned")
-            .is_some();
-
-        is_token_set || is_login_method_set
-    }
-
     #[allow(missing_docs)]
     #[cfg(feature = "internal")]
     pub fn get_kdf(&self) -> Result<Kdf, NotAuthenticatedError> {

crates/bitwarden-core/src/client/internal.rs

@@ -63,14 +63,10 @@ pub(crate) async fn get_user_api_key(
 }
 
 fn get_login_method(client: &Client) -> Result<Arc<LoginMethod>, NotAuthenticatedError> {
-    if client.internal.is_authed() {
-        client
-            .internal
-            .get_login_method()
-            .ok_or(NotAuthenticatedError)
-    } else {
-        Err(NotAuthenticatedError)
-    }
+    client
+        .internal
+        .get_login_method()
+        .ok_or(NotAuthenticatedError)
 }
 
 /// Build the secret verification request.

crates/bitwarden-core/src/platform/get_user_api_key.rs

@@ -1,5 +1,5 @@
 extern crate console_error_panic_hook;
-use std::fmt::Display;
+use std::{fmt::Display, sync::Arc};
 
 use bitwarden_core::{key_management::CryptoClient, Client, ClientSettings};
 use bitwarden_error::bitwarden_error;
@@ -8,7 +8,10 @@ use bitwarden_generators::GeneratorClientsExt;
 use bitwarden_vault::{VaultClient, VaultClientExt};
 use wasm_bindgen::prelude::*;
 
-use crate::platform::PlatformClient;
+use crate::platform::{
+    token_provider::{JsTokenProvider, WasmClientManagedTokens},
+    PlatformClient,
+};
 
 #[allow(missing_docs)]
 #[wasm_bindgen]
@@ -18,8 +21,9 @@ pub struct BitwardenClient(pub(crate) Client);
 impl BitwardenClient {
     #[allow(missing_docs)]
     #[wasm_bindgen(constructor)]
-    pub fn new(settings: Option<ClientSettings>) -> Self {
-        Self(Client::new(settings))
+    pub fn new(settings: Option<ClientSettings>, token_provider: JsTokenProvider) -> Self {
+        let tokens = Arc::new(WasmClientManagedTokens::new(token_provider));
+        Self(Client::new_with_client_tokens(settings, tokens))
     }
 
     /// Test method, echoes back the input

crates/bitwarden-wasm-internal/src/client.rs

@@ -3,6 +3,7 @@ use bitwarden_vault::Cipher;
 use wasm_bindgen::prelude::wasm_bindgen;
 
 mod repository;
+pub mod token_provider;
 
 #[wasm_bindgen]
 pub struct PlatformClient(Client);

crates/bitwarden-wasm-internal/src/platform/mod.rs

@@ -0,0 +1,44 @@
+use bitwarden_core::client::internal::ClientManagedTokens;
+use bitwarden_threading::ThreadBoundRunner;
+use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
+
+#[wasm_bindgen(typescript_custom_section)]
+const TOKEN_CUSTOM_TS_TYPE: &'static str = r#"
+export interface TokenProvider {
+    get_access_token(): Promise<string>;
+}
+"#;
+
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(js_name = TokenProvider)]
+    pub type JsTokenProvider;
+
+    #[wasm_bindgen(method)]
+    pub async fn get_access_token(this: &JsTokenProvider) -> JsValue;
+}
+
+/// Thread-bound runner for JavaScript token provider
+pub(crate) struct WasmClientManagedTokens(ThreadBoundRunner<JsTokenProvider>);
+
+impl WasmClientManagedTokens {
+    pub fn new(js_provider: JsTokenProvider) -> Self {
+        Self(ThreadBoundRunner::new(js_provider))
+    }
+}
+
+impl std::fmt::Debug for WasmClientManagedTokens {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WasmClientManagedTokens").finish()
+    }
+}
+
+#[async_trait::async_trait]
+impl ClientManagedTokens for WasmClientManagedTokens {
+    async fn get_access_token(&self) -> Option<String> {
+        self.0
+            .run_in_thread(|c| async move { c.get_access_token().await.as_string() })
+            .await
+            .unwrap_or_default()
+    }
+}