CX-6295 backend access token authorization wiring (#25236)

GitOrigin-RevId: 7832cb488329e89c9eab5dfaa10827ebc89daf42

Author: pashabitz

Cargo.lock

@@ -20,6 +20,7 @@ use std::{
 
 use anyhow::Context;
 use authentication::{
+    application_auth::ApplicationAuth,
     validate_id_token,
     Auth0IdToken,
 };
@@ -411,6 +412,7 @@ pub struct Application<RT: Runtime> {
     log_visibility: Arc<dyn LogVisibility<RT>>,
     module_cache: ModuleCache<RT>,
     system_env_var_names: HashSet<EnvVarName>,
+    app_auth: Arc<ApplicationAuth>,
 }
 
 impl<RT: Runtime> Clone for Application<RT> {
@@ -443,6 +445,7 @@ impl<RT: Runtime> Clone for Application<RT> {
             log_visibility: self.log_visibility.clone(),
             module_cache: self.module_cache.clone(),
             system_env_var_names: self.system_env_var_names.clone(),
+            app_auth: self.app_auth.clone(),
         }
     }
 }
@@ -472,6 +475,7 @@ impl<RT: Runtime> Application<RT> {
         log_visibility: Arc<dyn LogVisibility<RT>>,
         snapshot_import_pause_client: PauseClient,
         scheduled_jobs_pause_client: PauseClient,
+        app_auth: Arc<ApplicationAuth>,
     ) -> anyhow::Result<Self> {
         let module_cache =
             ModuleCache::new(runtime.clone(), database.clone(), modules_storage.clone()).await;
@@ -600,6 +604,7 @@ impl<RT: Runtime> Application<RT> {
             log_visibility,
             module_cache,
             system_env_var_names: system_env_vars.into_keys().collect(),
+            app_auth,
         })
     }
 
@@ -718,6 +723,10 @@ impl<RT: Runtime> Application<RT> {
         self.database.latest_snapshot()
     }
 
+    pub fn app_auth(&self) -> Arc<ApplicationAuth> {
+        self.app_auth.clone()
+    }
+
     pub async fn search_with_compiled_query(
         &self,
         index_id: IndexId,

crates/application/src/lib.rs

@@ -7,6 +7,10 @@ use std::{
 };
 
 use async_trait::async_trait;
+use authentication::{
+    access_token_auth::NullAccessTokenAuth,
+    application_auth::ApplicationAuth,
+};
 use cmd_util::env::config_test;
 use common::{
     bootstrap_model::index::database_index::IndexedFields,
@@ -238,6 +242,10 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
             Arc::new(AllowLogging),
             snapshot_import_pause_client,
             args.scheduled_jobs_pause_client,
+            Arc::new(ApplicationAuth::new(
+                kb.clone(),
+                Arc::new(NullAccessTokenAuth),
+            )),
         )
         .await?;
 

crates/application/src/test_helpers.rs

@@ -10,6 +10,7 @@ doctest = false
 
 [dependencies]
 anyhow = { workspace = true }
+async-trait = { workspace = true }
 base64 = { workspace = true }
 biscuit = { workspace = true }
 chrono = { workspace = true }

crates/authentication/Cargo.toml

@@ -0,0 +1,24 @@
+use async_trait::async_trait;
+use keybroker::Identity;
+
+/// Logic to check authorization based on Access Token
+#[async_trait]
+pub trait AccessTokenAuth: Send + Sync {
+    async fn is_authorized(
+        &self,
+        instance_name: &str,
+        access_token: &str,
+    ) -> anyhow::Result<Identity>;
+}
+pub struct NullAccessTokenAuth;
+
+#[async_trait]
+impl AccessTokenAuth for NullAccessTokenAuth {
+    async fn is_authorized(
+        &self,
+        _instance_name: &str,
+        _access_token: &str,
+    ) -> anyhow::Result<Identity> {
+        anyhow::bail!("Access token authorization is not supported")
+    }
+}

crates/authentication/src/access_token_auth.rs

@@ -0,0 +1,49 @@
+use std::sync::Arc;
+
+use keybroker::{
+    Identity,
+    KeyBroker,
+};
+
+use crate::access_token_auth::AccessTokenAuth;
+
+pub struct ApplicationAuth {
+    key_broker: KeyBroker,
+    access_token_auth: Arc<dyn AccessTokenAuth>,
+}
+
+// Encapsulates auth logic supporting both legacy Deploy Keys and new Convex
+// Access tokens
+impl ApplicationAuth {
+    pub fn new(key_broker: KeyBroker, access_token_auth: Arc<dyn AccessTokenAuth>) -> Self {
+        Self {
+            key_broker,
+            access_token_auth,
+        }
+    }
+
+    pub async fn check_key(
+        &self,
+        admin_key_or_access_token: String,
+        instance_name: String,
+    ) -> anyhow::Result<Identity> {
+        if admin_key_or_access_token.contains('|')
+            || self
+                .key_broker
+                .is_encrypted_admin_key(&admin_key_or_access_token)
+        {
+            // assume this is a legacy Deploy Key
+            // This is either a pipe-delimited deployment specific key
+            // or an encrypted admin key.
+            // The latter is used by smoke tests.
+            self.key_broker.check_admin_key(&admin_key_or_access_token)
+        } else {
+            // assume this is an Access Token
+            // Access Tokens are base64 encoded strings and do not have pipes
+            // in them
+            self.access_token_auth
+                .is_authorized(&instance_name, &admin_key_or_access_token)
+                .await
+        }
+    }
+}

crates/authentication/src/application_auth.rs

@@ -46,6 +46,9 @@ use serde::{
 use sync_types::AuthenticationToken;
 use url::Url;
 
+pub mod access_token_auth;
+pub mod application_auth;
+
 /// Issuer for API access tokens
 pub static CONVEX_AUTH_URL: LazyLock<Url> =
     LazyLock::new(|| Url::parse("https://auth.convex.dev/").unwrap());

crates/authentication/src/lib.rs

@@ -185,6 +185,7 @@ impl FetchClient for StaticFetchClient {
 
 pub enum InternalFetchPurpose {
     UsageTracking,
+    AccessTokenAuth,
 }
 
 #[cfg(test)]

crates/common/src/http/fetch.rs

@@ -480,6 +480,18 @@ impl AdminIdentity {
             key,
         })
     }
+
+    pub fn new_for_access_token(
+        instance_name: String,
+        member_id: MemberId,
+        access_token: String,
+    ) -> Self {
+        Self {
+            instance_name,
+            member_id,
+            key: access_token,
+        }
+    }
 }
 
 #[cfg(any(test, feature = "testing"))]
@@ -606,6 +618,16 @@ impl KeyBroker {
         )
     }
 
+    pub fn is_encrypted_admin_key(&self, key: &str) -> bool {
+        let (_, encrypted_part) = split_admin_key(key)
+            .map(|(name, key)| (Some(remove_type_prefix_from_instance_name(name)), key))
+            .unwrap_or((None, key));
+        let admin_key: Result<AdminKeyProto, _> = self
+            .encryptor
+            .decode_proto(ADMIN_KEY_VERSION, encrypted_part);
+        admin_key.is_ok()
+    }
+
     pub fn check_admin_key(&self, key: &str) -> anyhow::Result<Identity> {
         let (instance_name, encrypted_part) = split_admin_key(key)
             .map(|(name, key)| (Some(remove_type_prefix_from_instance_name(name)), key))

crates/keybroker/src/broker.rs

@@ -390,9 +390,11 @@ pub async fn push_config_handler(
         .into_iter()
         .map(|m| m.try_into())
         .collect::<anyhow::Result<Vec<_>>>()?;
+
     let identity = application
-        .key_broker()
-        .check_admin_key(&config.admin_key)
+        .app_auth()
+        .check_key(config.admin_key, application.instance_name())
+        .await
         .context("bad admin key error")?;
 
     let udf_server_version = Version::parse(&config.udf_server_version).context(