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 upstream_oauth_status(
45        &self,
46        upstream: &str,
47        subject: &str,
48    ) -> Result<UpstreamOauthStatusView, GatewayManagerError> {
49        let manager = self.require_oauth_manager(upstream)?;
50        let status = manager
51            .credential_status(subject)
52            .await
53            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))?;
54        let Some(status) = status else {
55            return Ok(UpstreamOauthStatusView {
56                authenticated: false,
57                upstream: upstream.to_owned(),
58                state: UpstreamOauthConnectionState::Disconnected,
59                access_token_expires_at: None,
60                seconds_until_expiry: None,
61                refresh_token_present: false,
62            });
63        };
64        let now = now_unix()?;
65        let seconds = status.access_token_expires_at.saturating_sub(now);
66        let state = if status.access_token_expires_at <= now {
67            UpstreamOauthConnectionState::Expired
68        } else if seconds <= TOKEN_EXPIRY_WARNING_SECS {
69            UpstreamOauthConnectionState::Expiring
70        } else {
71            UpstreamOauthConnectionState::Connected
72        };
73        Ok(UpstreamOauthStatusView {
74            authenticated: matches!(
75                state,
76                UpstreamOauthConnectionState::Connected | UpstreamOauthConnectionState::Expiring
77            ),
78            upstream: upstream.to_owned(),
79            state,
80            access_token_expires_at: Some(status.access_token_expires_at),
81            seconds_until_expiry: Some(seconds),
82            refresh_token_present: status.refresh_token_present,
83        })
84    }
85
86    pub async fn clear_upstream_credentials(
87        &self,
88        upstream: &str,
89        subject: &str,
90    ) -> Result<(), GatewayManagerError> {
91        let manager = self.require_oauth_manager(upstream)?;
92        manager
93            .clear_credentials(subject)
94            .await
95            .map_err(|error| GatewayManagerError::OAuth(error.to_string()))?;
96        if let Some(runtime) = self
97            .oauth_runtime
98            .read()
99            .expect("gateway oauth runtime poisoned")
100            .as_ref()
101        {
102            runtime.evict_subject(upstream, subject);
103            self.pool
104                .read()
105                .expect("gateway pool poisoned")
106                .evict_oauth_subject(upstream, subject);
107        }
108        Ok(())
109    }
110
111    fn require_oauth_manager(
112        &self,
113        upstream: &str,
114    ) -> Result<std::sync::Arc<dyn UpstreamOAuthManager>, GatewayManagerError> {
115        let runtime = self
116            .oauth_runtime
117            .read()
118            .expect("gateway oauth runtime poisoned")
119            .clone()
120            .ok_or_else(|| {
121                GatewayManagerError::OAuth("upstream OAuth runtime is not configured".to_owned())
122            })?;
123        runtime.manager(upstream).ok_or_else(|| {
124            GatewayManagerError::OAuth(format!("upstream `{upstream}` has no OAuth manager"))
125        })
126    }
127}
128
129fn now_unix() -> Result<i64, GatewayManagerError> {
130    SystemTime::now()
131        .duration_since(UNIX_EPOCH)
132        .map(|duration| duration.as_secs() as i64)
133        .map_err(|error| GatewayManagerError::OAuth(format!("system clock error: {error}")))
134}
135
136#[cfg(test)]
137#[path = "oauth_lifecycle_tests.rs"]
138mod tests;
soma_gateway/gateway/manager/oauth_lifecycle.rs

soma_gateway/gateway/manager/
oauth_lifecycle.rs
