soma_gateway/gateway/manager/
oauth_lifecycle.rs

1use std::time::{SystemTime, UNIX_EPOCH};
2
3use serde::Serialize;
4use soma_mcp_client::oauth::{BeginAuthorization, UpstreamOAuthManager};
5
6use super::{GatewayManager, GatewayManagerError};
7
8const TOKEN_EXPIRY_WARNING_SECS: i64 = 300;
9
10#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
11pub struct UpstreamOauthStatusView {
12    pub authenticated: bool,
13    pub upstream: String,
14    pub state: UpstreamOauthConnectionState,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub access_token_expires_at: Option<i64>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub seconds_until_expiry: Option<i64>,
19    #[serde(default)]
20    pub refresh_token_present: bool,
21}
22
23#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
24#[serde(rename_all = "snake_case")]
25pub enum UpstreamOauthConnectionState {
26    Connected,
27    Expiring,
28    Expired,
29    Disconnected,
30}
31
32impl GatewayManager {
33    pub async fn begin_upstream_authorization(
34        &self,
35        upstream: &str,
36        subject: &str,
37    ) -> Result<BeginAuthorization, GatewayManagerError> {
38        self.require_oauth_manager(upstream)?
39            .begin_authorization(subject)
40            .await
41            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))
42    }
43
44    pub async fn complete_upstream_authorization(
45        &self,
46        upstream: &str,
47        subject: &str,
48        code: &str,
49        state: &str,
50        issuer: Option<&str>,
51    ) -> Result<(), GatewayManagerError> {
52        let manager = self.require_oauth_manager(upstream)?;
53        manager
54            .complete_authorization(subject, code, state, issuer)
55            .await
56            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))?;
57        if let Some(runtime) = self
58            .oauth_runtime
59            .read()
60            .expect("gateway oauth runtime poisoned")
61            .as_ref()
62        {
63            runtime.evict_subject(upstream, subject);
64            self.pool
65                .read()
66                .expect("gateway pool poisoned")
67                .evict_oauth_subject(upstream, subject);
68        }
69        Ok(())
70    }
71
72    pub async fn upstream_oauth_status(
73        &self,
74        upstream: &str,
75        subject: &str,
76    ) -> Result<UpstreamOauthStatusView, GatewayManagerError> {
77        let manager = self.require_oauth_manager(upstream)?;
78        let status = manager
79            .credential_status(subject)
80            .await
81            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))?;
82        let Some(status) = status else {
83            return Ok(UpstreamOauthStatusView {
84                authenticated: false,
85                upstream: upstream.to_owned(),
86                state: UpstreamOauthConnectionState::Disconnected,
87                access_token_expires_at: None,
88                seconds_until_expiry: None,
89                refresh_token_present: false,
90            });
91        };
92        let now = now_unix()?;
93        let seconds = status.access_token_expires_at.saturating_sub(now);
94        let state = if status.access_token_expires_at <= now {
95            UpstreamOauthConnectionState::Expired
96        } else if seconds <= TOKEN_EXPIRY_WARNING_SECS {
97            UpstreamOauthConnectionState::Expiring
98        } else {
99            UpstreamOauthConnectionState::Connected
100        };
101        Ok(UpstreamOauthStatusView {
102            authenticated: matches!(
103                state,
104                UpstreamOauthConnectionState::Connected | UpstreamOauthConnectionState::Expiring
105            ),
106            upstream: upstream.to_owned(),
107            state,
108            access_token_expires_at: Some(status.access_token_expires_at),
109            seconds_until_expiry: Some(seconds),
110            refresh_token_present: status.refresh_token_present,
111        })
112    }
113
114    pub async fn clear_upstream_credentials(
115        &self,
116        upstream: &str,
117        subject: &str,
118    ) -> Result<(), GatewayManagerError> {
119        let manager = self.require_oauth_manager(upstream)?;
120        manager
121            .clear_credentials(subject)
122            .await
123            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))?;
124        if let Some(runtime) = self
125            .oauth_runtime
126            .read()
127            .expect("gateway oauth runtime poisoned")
128            .as_ref()
129        {
130            runtime.evict_subject(upstream, subject);
131            self.pool
132                .read()
133                .expect("gateway pool poisoned")
134                .evict_oauth_subject(upstream, subject);
135        }
136        Ok(())
137    }
138
139    fn require_oauth_manager(
140        &self,
141        upstream: &str,
142    ) -> Result<std::sync::Arc<dyn UpstreamOAuthManager>, GatewayManagerError> {
143        let runtime = self
144            .oauth_runtime
145            .read()
146            .expect("gateway oauth runtime poisoned")
147            .clone()
148            .ok_or_else(|| {
149                GatewayManagerError::OAuth("upstream OAuth runtime is not configured".to_owned())
150            })?;
151        runtime.manager(upstream).ok_or_else(|| {
152            GatewayManagerError::OAuth(format!("upstream `{upstream}` has no OAuth manager"))
153        })
154    }
155}
156
157fn now_unix() -> Result<i64, GatewayManagerError> {
158    SystemTime::now()
159        .duration_since(UNIX_EPOCH)
160        .map(|duration| duration.as_secs() as i64)
161        .map_err(|error| GatewayManagerError::OAuth(format!("system clock error: {error}")))
162}
163
164#[cfg(test)]
165#[path = "oauth_lifecycle_tests.rs"]
166mod tests;
soma_gateway/gateway/manager/oauth_lifecycle.rs

soma_gateway/gateway/manager/
oauth_lifecycle.rs
