systemprompt_cloud/oauth/
client.rs

1//! Browser-driven OAuth login flow.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use std::sync::Arc;
7
8use axum::Router;
9use axum::extract::{Query, State};
10use axum::response::Html;
11use axum::routing::get;
12use reqwest::Client;
13use systemprompt_logging::CliService;
14use systemprompt_models::net::{HTTP_CONNECT_TIMEOUT, HTTP_DEFAULT_TIMEOUT};
15use tokio::sync::{Mutex, oneshot};
16
17use crate::OAuthProvider;
18use crate::constants::oauth::{CALLBACK_PORT, CALLBACK_TIMEOUT_SECS};
19use crate::error::{CloudError, CloudResult};
20
21#[derive(serde::Deserialize)]
22struct CallbackParams {
23    access_token: Option<String>,
24    error: Option<String>,
25    error_description: Option<String>,
26}
27
28#[derive(serde::Deserialize)]
29struct AuthorizeResponse {
30    authorize_url: String,
31}
32
33#[derive(Debug, Clone, Copy)]
34pub struct OAuthTemplates {
35    pub success_html: &'static str,
36    pub error_html: &'static str,
37}
38
39struct CallbackState {
40    tx: Mutex<Option<oneshot::Sender<CloudResult<String>>>>,
41    success_html: String,
42    error_html: String,
43}
44
45async fn callback_handler(
46    State(state): State<Arc<CallbackState>>,
47    Query(params): Query<CallbackParams>,
48) -> Html<String> {
49    let result: CloudResult<String> = if let Some(error) = params.error {
50        let desc = params
51            .error_description
52            .unwrap_or_else(|| "(no description provided)".into());
53        Err(CloudError::OAuthFlow {
54            message: format!("OAuth error: {error} - {desc}"),
55        })
56    } else if let Some(token) = params.access_token {
57        Ok(token)
58    } else {
59        Err(CloudError::OAuthFlow {
60            message: "No token received in callback".to_owned(),
61        })
62    };
63
64    let sender = state.tx.lock().await.take();
65    let Some(sender) = sender else {
66        return Html(state.error_html.clone());
67    };
68
69    let is_success = result.is_ok();
70    if sender.send(result).is_err() {
71        tracing::warn!("OAuth result receiver dropped before result could be sent");
72    }
73
74    if is_success {
75        Html(state.success_html.clone())
76    } else {
77        Html(state.error_html.clone())
78    }
79}
80
81async fn fetch_authorize_url(
82    api_url: &str,
83    provider: OAuthProvider,
84    redirect_uri: &str,
85) -> CloudResult<String> {
86    let client = Client::builder()
87        .connect_timeout(HTTP_CONNECT_TIMEOUT)
88        .timeout(HTTP_DEFAULT_TIMEOUT)
89        .build()?;
90    let oauth_endpoint = format!(
91        "{}/api/v1/auth/oauth/{}?redirect_uri={}",
92        api_url,
93        provider.as_str(),
94        urlencoding::encode(redirect_uri)
95    );
96
97    let response = client.get(&oauth_endpoint).send().await?;
98
99    if !response.status().is_success() {
100        let status = response.status();
101        let body = response.text().await.unwrap_or_else(|e| {
102            tracing::warn!(error = %e, "Failed to read OAuth error response body");
103            format!("(body unreadable: {e})")
104        });
105        return Err(CloudError::OAuthFlow {
106            message: format!("Failed to get authorization URL ({status}): {body}"),
107        });
108    }
109
110    let auth_response: AuthorizeResponse = response.json().await?;
111    Ok(auth_response.authorize_url)
112}
113
114async fn await_callback(
115    listener: tokio::net::TcpListener,
116    app: Router,
117    rx: oneshot::Receiver<CloudResult<String>>,
118) -> CloudResult<String> {
119    let server = axum::serve(listener, app);
120
121    tokio::select! {
122        result = rx => {
123            result.map_err(|_e| CloudError::OAuthFlow { message: "Authentication cancelled".to_owned() })?
124        }
125        _ = server => {
126            Err(CloudError::OAuthFlow { message: "Server stopped unexpectedly".to_owned() })
127        }
128        () = tokio::time::sleep(std::time::Duration::from_secs(CALLBACK_TIMEOUT_SECS)) => {
129            Err(CloudError::OAuthFlow { message: format!("Authentication timed out after {CALLBACK_TIMEOUT_SECS} seconds") })
130        }
131    }
132}
133
134pub async fn run_oauth_flow(
135    api_url: &str,
136    provider: OAuthProvider,
137    templates: OAuthTemplates,
138) -> CloudResult<String> {
139    let (tx, rx) = oneshot::channel::<CloudResult<String>>();
140    let state = Arc::new(CallbackState {
141        tx: Mutex::new(Some(tx)),
142        success_html: templates.success_html.to_owned(),
143        error_html: templates.error_html.to_owned(),
144    });
145
146    let app = Router::new()
147        .route("/callback", get(callback_handler))
148        .with_state(state);
149    let addr = format!("127.0.0.1:{CALLBACK_PORT}");
150    let listener = crate::callback_listener::bind_callback_listener(CALLBACK_PORT)?;
151
152    CliService::info(&format!("Starting authentication server on http://{addr}"));
153
154    let redirect_uri = format!("http://127.0.0.1:{CALLBACK_PORT}/callback");
155
156    CliService::info("Fetching authorization URL...");
157    let auth_url = fetch_authorize_url(api_url, provider, &redirect_uri).await?;
158
159    CliService::info(&format!(
160        "Opening browser for {} authentication...",
161        provider.display_name()
162    ));
163    CliService::info(&format!("URL: {auth_url}"));
164
165    if let Err(e) = open::that(&auth_url) {
166        CliService::warning(&format!("Could not open browser automatically: {e}"));
167        CliService::info("Please open this URL manually:");
168        CliService::key_value("URL", &auth_url);
169    }
170
171    CliService::info("Waiting for authentication...");
172    CliService::info(&format!("(timeout in {CALLBACK_TIMEOUT_SECS} seconds)"));
173
174    await_callback(listener, app, rx).await
175}
systemprompt_cloud/oauth/client.rs

systemprompt_cloud/oauth/
client.rs
