Skip to main content

skiff_cli/oauth/
setup.rs

1//! Build and run OAuth flows via `rmcp::transport::auth`.
2
3use std::time::Duration;
4
5use rmcp::transport::auth::{
6    AuthClient, AuthorizationRequest, ClientCredentialsConfig, CredentialStore, OAuthState,
7};
8
9use crate::cli::args::GlobalArgs;
10use crate::coerce::resolve_secret;
11use crate::error::{Error, Result};
12use crate::oauth::browser::open_authorization_url;
13use crate::oauth::callback::wait_for_callback;
14use crate::oauth::config::{parse_oauth_flow, resolve_redirect_uri, OAuthFlow, OAuthOptions};
15use crate::oauth::store::{clear_server_credentials, FileCredentialStore};
16
17const CALLBACK_TIMEOUT: Duration = Duration::from_secs(300);
18
19/// Ready-to-use OAuth session wrapping rmcp `AuthClient`.
20pub struct OAuthReady {
21    auth_client: AuthClient<reqwest::Client>,
22}
23
24impl OAuthReady {
25    pub fn auth_client(&self) -> &AuthClient<reqwest::Client> {
26        &self.auth_client
27    }
28
29    pub fn into_auth_client(self) -> AuthClient<reqwest::Client> {
30        self.auth_client
31    }
32
33    pub async fn access_token(&self) -> Result<String> {
34        self.auth_client
35            .get_access_token()
36            .await
37            .map_err(|e| Error::runtime(format!("OAuth get_access_token: {e}")))
38    }
39}
40
41/// Whether any OAuth flag requests auth.
42pub fn oauth_wanted(args: &GlobalArgs) -> bool {
43    args.oauth || args.oauth_client_id.is_some() || args.oauth_client_secret.is_some()
44}
45
46/// Validate CLI oauth flags and resolve secrets into [`OAuthOptions`].
47pub fn options_from_args(args: &GlobalArgs) -> Result<Option<OAuthOptions>> {
48    if !oauth_wanted(args) {
49        return Ok(None);
50    }
51    if args.mcp_stdio.is_some() {
52        return Err(Error::usage(
53            "OAuth is not supported with --mcp-stdio (HTTP discovery required)",
54        ));
55    }
56    if args.oauth_client_secret.is_some() && args.oauth_client_id.is_none() {
57        return Err(Error::usage(
58            "--oauth-client-secret requires --oauth-client-id",
59        ));
60    }
61
62    let flow = parse_oauth_flow(&args.oauth_flow)?;
63    let client_id = args
64        .oauth_client_id
65        .as_deref()
66        .map(resolve_secret)
67        .transpose()?;
68    let client_secret = args
69        .oauth_client_secret
70        .as_deref()
71        .map(resolve_secret)
72        .transpose()?;
73
74    if flow == OAuthFlow::ClientCredentials && (client_id.is_none() || client_secret.is_none()) {
75        return Err(Error::usage(
76            "--oauth-flow client_credentials requires both --oauth-client-id and --oauth-client-secret",
77        ));
78    }
79
80    if let Some(uri) = &args.oauth_redirect_uri {
81        crate::oauth::config::validate_redirect_uri(uri)?;
82    }
83
84    Ok(Some(OAuthOptions {
85        client_id,
86        client_secret,
87        client_name: args.oauth_client_name.clone(),
88        scope: args.oauth_scope.clone(),
89        redirect_uri: args.oauth_redirect_uri.clone(),
90        flow,
91    }))
92}
93
94/// Discoverable server URL for OAuth metadata (MCP URL, or OpenAPI HTTP base).
95pub fn discovery_url_from_args(args: &GlobalArgs) -> Result<String> {
96    if let Some(url) = &args.mcp {
97        return Ok(url.clone());
98    }
99    if let Some(gql) = &args.graphql {
100        return Ok(gql.clone());
101    }
102    if let Some(spec) = &args.spec {
103        if spec.starts_with("http://") || spec.starts_with("https://") {
104            return Ok(spec.clone());
105        }
106        if let Some(base) = &args.base_url {
107            return Ok(base.clone());
108        }
109        return Err(Error::usage(
110            "OAuth with a local --spec requires --base-url (or an HTTP --spec URL)",
111        ));
112    }
113    Err(Error::usage(
114        "OAuth requires --mcp, --graphql, or an HTTP --spec/--base-url",
115    ))
116}
117
118/// Run OAuth (or restore cached tokens) and return an authorized client.
119pub async fn authorize(server_url: &str, opts: &OAuthOptions) -> Result<OAuthReady> {
120    let store = FileCredentialStore::for_server(server_url);
121    let oauth_http = reqwest::Client::builder()
122        .timeout(Duration::from_secs(30))
123        .build()
124        .map_err(|e| Error::runtime(format!("oauth http client: {e}")))?;
125
126    let mut state = OAuthState::new(server_url, Some(oauth_http))
127        .await
128        .map_err(|e| Error::runtime(format!("OAuth init: {e}")))?;
129
130    attach_store(&mut state, store.clone())?;
131
132    if try_restore(&mut state, server_url).await? {
133        return finish(state);
134    }
135
136    if opts.use_client_credentials() {
137        let client_id = opts
138            .client_id
139            .clone()
140            .ok_or_else(|| Error::usage("client credentials requires --oauth-client-id"))?;
141        let client_secret = opts
142            .client_secret
143            .clone()
144            .ok_or_else(|| Error::usage("client credentials requires --oauth-client-secret"))?;
145        let scopes = scopes_vec(opts);
146        let config = ClientCredentialsConfig::ClientSecret {
147            client_id,
148            client_secret,
149            scopes,
150            resource: Some(server_url.to_string()),
151        };
152        state
153            .authenticate_client_credentials(config)
154            .await
155            .map_err(|e| Error::runtime(format!("client credentials auth failed: {e}")))?;
156        return finish(state);
157    }
158
159    // Authorization code + PKCE
160    let sticky = store.load_sticky_redirect_uri();
161    let redirect = resolve_redirect_uri(opts.redirect_uri.as_deref(), sticky.as_deref())?;
162    if sticky.as_deref() != Some(redirect.as_str()) {
163        let _ = store.clear().await;
164    }
165    store.save_sticky_redirect_uri(&redirect)?;
166
167    let mut request =
168        AuthorizationRequest::new(redirect.clone()).with_client_name(opts.client_name.clone());
169    if let Some(id) = &opts.client_id {
170        request = request.with_preregistered_client(id.clone());
171        if let Some(sec) = &opts.client_secret {
172            request = request.with_client_secret(sec.clone());
173        }
174    }
175    if let Some(scope) = &opts.scope {
176        request = request.with_scopes(scope.split_whitespace().map(|s| s.to_string()));
177    }
178
179    let redirect_for_cb = redirect.clone();
180    let cb_handle =
181        std::thread::spawn(move || wait_for_callback(&redirect_for_cb, CALLBACK_TIMEOUT));
182
183    state
184        .start_authorization(request)
185        .await
186        .map_err(|e| Error::runtime(format!("start authorization failed: {e}")))?;
187
188    let auth_url = state
189        .get_authorization_url()
190        .await
191        .map_err(|e| Error::runtime(format!("get authorization URL: {e}")))?;
192    open_authorization_url(&auth_url);
193
194    let cb = cb_handle
195        .join()
196        .map_err(|_| Error::runtime("callback thread panicked"))??;
197    let code = cb.code.as_deref().unwrap_or("");
198    let csrf = cb.state.as_deref().unwrap_or("");
199    state
200        .handle_callback_with_issuer(code, csrf, cb.iss.as_deref())
201        .await
202        .map_err(|e| Error::runtime(format!("OAuth callback exchange failed: {e}")))?;
203
204    finish(state)
205}
206
207fn attach_store(state: &mut OAuthState, store: FileCredentialStore) -> Result<()> {
208    match state {
209        OAuthState::Unauthorized(manager) => {
210            manager.set_credential_store(store);
211            Ok(())
212        }
213        _ => Err(Error::runtime(
214            "OAuth state not Unauthorized at store attach",
215        )),
216    }
217}
218
219async fn try_restore(state: &mut OAuthState, server_url: &str) -> Result<bool> {
220    match state {
221        OAuthState::Unauthorized(manager) => {
222            let ok = manager
223                .initialize_from_store()
224                .await
225                .map_err(|e| Error::runtime(format!("restore OAuth credentials: {e}")))?;
226            if !ok {
227                return Ok(false);
228            }
229            match manager.get_access_token().await {
230                Ok(_) => Ok(true),
231                Err(e) => {
232                    tracing::debug!("cached OAuth token unusable: {e}");
233                    let _ = FileCredentialStore::for_server(server_url).clear().await;
234                    Ok(false)
235                }
236            }
237        }
238        OAuthState::Authorized(_) => Ok(true),
239        _ => Ok(false),
240    }
241}
242
243fn finish(state: OAuthState) -> Result<OAuthReady> {
244    let manager = match state {
245        OAuthState::Authorized(m) | OAuthState::Unauthorized(m) => m,
246        OAuthState::Session(_) => {
247            return Err(Error::runtime("OAuth still in session state after auth"));
248        }
249        OAuthState::AuthorizedHttpClient(_) => {
250            return Err(Error::runtime("unexpected AuthorizedHttpClient state"));
251        }
252        _ => {
253            return Err(Error::runtime("unexpected OAuth state"));
254        }
255    };
256    let auth_client = AuthClient::new(reqwest::Client::default(), manager);
257    Ok(OAuthReady { auth_client })
258}
259
260fn scopes_vec(opts: &OAuthOptions) -> Vec<String> {
261    opts.scope
262        .as_deref()
263        .map(|s| s.split_whitespace().map(str::to_string).collect())
264        .unwrap_or_default()
265}
266
267pub fn clear_oauth_credentials(server_url: &str) -> Result<()> {
268    clear_server_credentials(server_url)
269}
270
271/// Convenience: setup from GlobalArgs if oauth wanted.
272pub async fn setup_from_args(args: &GlobalArgs) -> Result<Option<OAuthReady>> {
273    let Some(opts) = options_from_args(args)? else {
274        return Ok(None);
275    };
276    let url = discovery_url_from_args(args)?;
277    Ok(Some(authorize(&url, &opts).await?))
278}