Skip to main content

soma_integrations/
gateway_auth.rs

1//! Gateway-to-auth bridge: translates `soma-auth`'s upstream OAuth runtime
2//! (Soma's product OAuth storage, refresh, and credential cache) into
3//! `soma-gateway`'s generic `UpstreamOAuthProvider`/`UpstreamOAuthManager`
4//! traits, so the standalone gateway engine can authenticate to upstream MCP
5//! servers without knowing about Soma's auth types (plan section 3.20).
6//!
7//! Moved out of `apps/soma` (formerly `gateway_auth.rs`), where it was kept
8//! as a temporary bridge for mergeability — this crate is its permanent home
9//! per PR 11's acceptance criterion that `apps/soma` constructs adapters but
10//! contains none of their implementation logic.
11
12use std::{collections::BTreeMap, sync::Arc};
13
14use anyhow::{Context, Result};
15use futures::future::BoxFuture;
16use mcp_client::{
17    config::{GatewayUpstreamOauthMode, GatewayUpstreamOauthRegistration, UpstreamConfig},
18    oauth::{
19        BeginAuthorization, UpstreamOAuthCredentialStatus, UpstreamOAuthError,
20        UpstreamOAuthHttpClient, UpstreamOAuthManager, UpstreamOAuthProvider, UpstreamOAuthRuntime,
21    },
22    upstream::http_body_cap::BodyCappedHttpClient,
23};
24
25#[derive(Clone)]
26struct SomaOAuthProvider {
27    cache: soma_auth::upstream::cache::OauthClientCache,
28}
29
30impl UpstreamOAuthProvider for SomaOAuthProvider {
31    fn authenticated_http_client<'a>(
32        &'a self,
33        upstream: &'a UpstreamConfig,
34        subject: &'a str,
35        http_client: BodyCappedHttpClient,
36    ) -> BoxFuture<'a, Result<UpstreamOAuthHttpClient, UpstreamOAuthError>> {
37        Box::pin(async move {
38            let upstream = to_auth_upstream(upstream)?;
39            self.cache
40                .get_or_build_capped(&upstream, subject, http_client)
41                .await
42                .map_err(map_oauth_error)
43        })
44    }
45
46    fn evict_subject(&self, upstream: &str, subject: &str) {
47        self.cache.evict_subject(upstream, subject);
48    }
49
50    fn evict_upstream(&self, upstream: &str) {
51        self.cache.evict_upstream(upstream);
52    }
53}
54
55#[derive(Clone)]
56struct SomaOAuthManager {
57    manager: soma_auth::upstream::manager::UpstreamOauthManager,
58    cache: soma_auth::upstream::cache::OauthClientCache,
59}
60
61impl UpstreamOAuthManager for SomaOAuthManager {
62    fn begin_authorization<'a>(
63        &'a self,
64        subject: &'a str,
65    ) -> BoxFuture<'a, Result<BeginAuthorization, UpstreamOAuthError>> {
66        Box::pin(async move {
67            self.manager
68                .begin_authorization(subject)
69                .await
70                .map(|result| BeginAuthorization {
71                    authorization_url: result.authorization_url,
72                })
73                .map_err(map_oauth_error)
74        })
75    }
76
77    fn credential_status<'a>(
78        &'a self,
79        subject: &'a str,
80    ) -> BoxFuture<'a, Result<Option<UpstreamOAuthCredentialStatus>, UpstreamOAuthError>> {
81        Box::pin(async move {
82            self.manager
83                .credential_row(subject)
84                .await
85                .map(|row| {
86                    row.map(|row| UpstreamOAuthCredentialStatus {
87                        access_token_expires_at: row.access_token_expires_at,
88                        refresh_token_present: row.refresh_token_present,
89                    })
90                })
91                .map_err(map_oauth_error)
92        })
93    }
94
95    fn clear_credentials<'a>(
96        &'a self,
97        subject: &'a str,
98    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>> {
99        Box::pin(async move {
100            self.manager
101                .clear_credentials(subject)
102                .await
103                .map_err(map_oauth_error)?;
104            self.cache
105                .evict_subject(&self.manager.upstream_config().name, subject);
106            Ok(())
107        })
108    }
109
110    fn access_token<'a>(
111        &'a self,
112        subject: &'a str,
113    ) -> BoxFuture<'a, Result<String, UpstreamOAuthError>> {
114        Box::pin(async move {
115            let client = self
116                .cache
117                .get_or_build(self.manager.upstream_config(), subject)
118                .await
119                .map_err(map_oauth_error)?;
120            client
121                .get_access_token()
122                .await
123                .map_err(|error| UpstreamOAuthError::internal(error.to_string()))
124        })
125    }
126}
127
128pub async fn build_runtime(
129    upstreams: &[UpstreamConfig],
130    auth_config: &soma_auth::config::AuthConfig,
131    encryption_key: Option<&str>,
132) -> Result<Option<UpstreamOAuthRuntime>> {
133    let auth_upstreams = upstreams
134        .iter()
135        .filter(|upstream| upstream.oauth.is_some())
136        .map(to_auth_upstream)
137        .collect::<Result<Vec<_>, _>>()?;
138    let Some(runtime) = soma_auth::upstream::runtime::build_upstream_oauth_runtime(
139        &auth_upstreams,
140        auth_config,
141        encryption_key,
142    )
143    .await
144    .context("build Soma upstream OAuth runtime")?
145    else {
146        return Ok(None);
147    };
148
149    let provider = Arc::new(SomaOAuthProvider {
150        cache: runtime.cache.clone(),
151    });
152    let managers = runtime
153        .managers
154        .iter()
155        .map(|entry| {
156            (
157                entry.key().clone(),
158                Arc::new(SomaOAuthManager {
159                    manager: entry.value().clone(),
160                    cache: runtime.cache.clone(),
161                }) as Arc<dyn UpstreamOAuthManager>,
162            )
163        })
164        .collect::<BTreeMap<_, _>>();
165    Ok(Some(UpstreamOAuthRuntime::new(provider, managers)))
166}
167
168fn to_auth_upstream(
169    upstream: &UpstreamConfig,
170) -> Result<soma_auth::upstream::config::UpstreamConfig, UpstreamOAuthError> {
171    let oauth = upstream
172        .oauth
173        .as_ref()
174        .ok_or_else(|| UpstreamOAuthError::internal("upstream OAuth config is missing"))?;
175    let url = upstream
176        .url
177        .clone()
178        .ok_or_else(|| UpstreamOAuthError::internal("upstream OAuth URL is missing"))?;
179    Ok(soma_auth::upstream::config::UpstreamConfig {
180        name: upstream.name.clone(),
181        url: Some(url),
182        oauth: Some(soma_auth::upstream::config::UpstreamOauthConfig {
183            mode: match oauth.mode {
184                GatewayUpstreamOauthMode::AuthorizationCodePkce => {
185                    soma_auth::upstream::config::UpstreamOauthMode::AuthorizationCodePkce
186                }
187            },
188            registration: match &oauth.registration {
189                GatewayUpstreamOauthRegistration::ClientMetadataDocument { url } => {
190                    soma_auth::upstream::config::UpstreamOauthRegistration::ClientMetadataDocument {
191                        url: url.clone(),
192                    }
193                }
194                GatewayUpstreamOauthRegistration::Preregistered {
195                    client_id,
196                    client_secret_env,
197                } => soma_auth::upstream::config::UpstreamOauthRegistration::Preregistered {
198                    client_id: client_id.clone(),
199                    client_secret_env: client_secret_env.clone(),
200                },
201                GatewayUpstreamOauthRegistration::Dynamic => {
202                    soma_auth::upstream::config::UpstreamOauthRegistration::Dynamic
203                }
204            },
205            scopes: oauth.scopes.clone(),
206            prefer_client_metadata_document: oauth.prefer_client_metadata_document,
207        }),
208    })
209}
210
211fn map_oauth_error(error: soma_auth::upstream::types::OauthError) -> UpstreamOAuthError {
212    UpstreamOAuthError::new(error.kind(), error.to_string())
213}
214
215#[cfg(test)]
216#[path = "gateway_auth_tests.rs"]
217mod tests;