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 complete_authorization<'a>(
78        &'a self,
79        subject: &'a str,
80        code: &'a str,
81        state: &'a str,
82        issuer: Option<&'a str>,
83    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>> {
84        Box::pin(async move {
85            self.manager
86                .complete_authorization_callback(subject, code, state, issuer)
87                .await
88                .map_err(map_oauth_error)?;
89            self.cache
90                .evict_subject(&self.manager.upstream_config().name, subject);
91            Ok(())
92        })
93    }
94
95    fn credential_status<'a>(
96        &'a self,
97        subject: &'a str,
98    ) -> BoxFuture<'a, Result<Option<UpstreamOAuthCredentialStatus>, UpstreamOAuthError>> {
99        Box::pin(async move {
100            self.manager
101                .credential_row(subject)
102                .await
103                .map(|row| {
104                    row.map(|row| UpstreamOAuthCredentialStatus {
105                        access_token_expires_at: row.access_token_expires_at,
106                        refresh_token_present: row.refresh_token_present,
107                    })
108                })
109                .map_err(map_oauth_error)
110        })
111    }
112
113    fn clear_credentials<'a>(
114        &'a self,
115        subject: &'a str,
116    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>> {
117        Box::pin(async move {
118            self.manager
119                .clear_credentials(subject)
120                .await
121                .map_err(map_oauth_error)?;
122            self.cache
123                .evict_subject(&self.manager.upstream_config().name, subject);
124            Ok(())
125        })
126    }
127
128    fn access_token<'a>(
129        &'a self,
130        subject: &'a str,
131    ) -> BoxFuture<'a, Result<String, UpstreamOAuthError>> {
132        Box::pin(async move {
133            let client = self
134                .cache
135                .get_or_build(self.manager.upstream_config(), subject)
136                .await
137                .map_err(map_oauth_error)?;
138            client
139                .get_access_token()
140                .await
141                .map_err(|error| UpstreamOAuthError::internal(error.to_string()))
142        })
143    }
144}
145
146pub async fn build_runtime(
147    upstreams: &[UpstreamConfig],
148    auth_config: &soma_auth::config::AuthConfig,
149    encryption_key: Option<&str>,
150) -> Result<Option<UpstreamOAuthRuntime>> {
151    let auth_upstreams = upstreams
152        .iter()
153        .filter(|upstream| upstream.oauth.is_some())
154        .map(to_auth_upstream)
155        .collect::<Result<Vec<_>, _>>()?;
156    let Some(runtime) = soma_auth::upstream::runtime::build_upstream_oauth_runtime(
157        &auth_upstreams,
158        auth_config,
159        encryption_key,
160    )
161    .await
162    .context("build Soma upstream OAuth runtime")?
163    else {
164        return Ok(None);
165    };
166
167    let provider = Arc::new(SomaOAuthProvider {
168        cache: runtime.cache.clone(),
169    });
170    let managers = runtime
171        .managers
172        .iter()
173        .map(|entry| {
174            (
175                entry.key().clone(),
176                Arc::new(SomaOAuthManager {
177                    manager: entry.value().clone(),
178                    cache: runtime.cache.clone(),
179                }) as Arc<dyn UpstreamOAuthManager>,
180            )
181        })
182        .collect::<BTreeMap<_, _>>();
183    Ok(Some(UpstreamOAuthRuntime::new(provider, managers)))
184}
185
186fn to_auth_upstream(
187    upstream: &UpstreamConfig,
188) -> Result<soma_auth::upstream::config::UpstreamConfig, UpstreamOAuthError> {
189    let oauth = upstream
190        .oauth
191        .as_ref()
192        .ok_or_else(|| UpstreamOAuthError::internal("upstream OAuth config is missing"))?;
193    let url = upstream
194        .url
195        .clone()
196        .ok_or_else(|| UpstreamOAuthError::internal("upstream OAuth URL is missing"))?;
197    Ok(soma_auth::upstream::config::UpstreamConfig {
198        name: upstream.name.clone(),
199        url: Some(url),
200        oauth: Some(soma_auth::upstream::config::UpstreamOauthConfig {
201            mode: match oauth.mode {
202                GatewayUpstreamOauthMode::AuthorizationCodePkce => {
203                    soma_auth::upstream::config::UpstreamOauthMode::AuthorizationCodePkce
204                }
205            },
206            registration: match &oauth.registration {
207                GatewayUpstreamOauthRegistration::Auto => {
208                    soma_auth::upstream::config::UpstreamOauthRegistration::Auto
209                }
210                GatewayUpstreamOauthRegistration::ClientMetadataDocument { url } => {
211                    soma_auth::upstream::config::UpstreamOauthRegistration::ClientMetadataDocument {
212                        url: url.clone(),
213                    }
214                }
215                GatewayUpstreamOauthRegistration::Preregistered {
216                    client_id,
217                    client_secret_env,
218                } => soma_auth::upstream::config::UpstreamOauthRegistration::Preregistered {
219                    client_id: client_id.clone(),
220                    client_secret_env: client_secret_env.clone(),
221                },
222                GatewayUpstreamOauthRegistration::Dynamic => {
223                    soma_auth::upstream::config::UpstreamOauthRegistration::Dynamic
224                }
225            },
226            scopes: oauth.scopes.clone(),
227            prefer_client_metadata_document: oauth.prefer_client_metadata_document,
228        }),
229    })
230}
231
232fn map_oauth_error(error: soma_auth::upstream::types::OauthError) -> UpstreamOAuthError {
233    UpstreamOAuthError::new(error.kind(), error.to_string())
234}
235
236#[cfg(test)]
237#[path = "gateway_auth_tests.rs"]
238mod tests;