1use 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;