soma_integrations/
gateway_auth.rs1use 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;