systemprompt_cloud/oauth/
client.rs

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

systemprompt_cloud/oauth/
client.rs
