systemprompt_cloud/operator/
login.rs

1//! Browser authorization-code flow with PKCE against an OIDC issuer.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use std::sync::Arc;
7use std::time::Duration;
8
9use axum::Router;
10use axum::extract::{Query, State};
11use axum::response::Html;
12use axum::routing::get;
13use base64::prelude::*;
14use chrono::Utc;
15use serde::Deserialize;
16use sha2::{Digest, Sha256};
17use systemprompt_identifiers::{ClientId, CloudAuthToken, Email};
18use systemprompt_logging::CliService;
19use systemprompt_models::net::{HTTP_CONNECT_TIMEOUT, HTTP_DEFAULT_TIMEOUT};
20use tokio::sync::{Mutex, oneshot};
21use uuid::Uuid;
22
23use super::endpoints::{AuthorizeParams, CodeRedemption, authorize_url, discover, redeem_code};
24use super::{OperatorCredentials, identity_claims};
25use crate::OAuthTemplates;
26use crate::constants::oauth::{CALLBACK_PORT, CALLBACK_TIMEOUT_SECS};
27use crate::error::{CloudError, CloudResult};
28
29#[derive(Debug, Clone)]
30pub struct OperatorLoginRequest {
31    pub issuer: String,
32    pub client_id: ClientId,
33    pub scopes: Vec<String>,
34}
35
36pub struct LoginIo<'a> {
37    pub callback_port: u16,
38    pub open_browser: &'a (dyn Fn(&str) -> std::io::Result<()> + Sync),
39}
40
41impl std::fmt::Debug for LoginIo<'_> {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.debug_struct("LoginIo")
44            .field("callback_port", &self.callback_port)
45            .finish_non_exhaustive()
46    }
47}
48
49#[derive(Deserialize)]
50struct CallbackParams {
51    code: Option<String>,
52    state: Option<String>,
53    error: Option<String>,
54    error_description: Option<String>,
55}
56
57struct CallbackState {
58    expected_state: String,
59    tx: Mutex<Option<oneshot::Sender<CloudResult<String>>>>,
60    success_html: &'static str,
61    error_html: &'static str,
62}
63
64fn random_token() -> String {
65    format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple())
66}
67
68fn pkce_challenge(verifier: &str) -> String {
69    BASE64_URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()))
70}
71
72async fn callback_handler(
73    State(state): State<Arc<CallbackState>>,
74    Query(params): Query<CallbackParams>,
75) -> Html<&'static str> {
76    let result = if let Some(error) = params.error {
77        Err(CloudError::operator_auth(format!(
78            "{error}: {}",
79            params.error_description.unwrap_or_default()
80        )))
81    } else if params.state.as_deref() != Some(state.expected_state.as_str()) {
82        Err(CloudError::operator_auth("callback state did not match"))
83    } else {
84        params
85            .code
86            .ok_or_else(|| CloudError::operator_auth("callback carried no authorization code"))
87    };
88    let is_success = result.is_ok();
89    let sender = state.tx.lock().await.take();
90    if let Some(sender) = sender
91        && sender.send(result).is_err()
92    {
93        tracing::warn!("Operator login receiver dropped before the callback arrived");
94    }
95    Html(if is_success {
96        state.success_html
97    } else {
98        state.error_html
99    })
100}
101
102async fn await_code(
103    listener: tokio::net::TcpListener,
104    app: Router,
105    rx: oneshot::Receiver<CloudResult<String>>,
106) -> CloudResult<String> {
107    let server = axum::serve(listener, app);
108    tokio::select! {
109        result = rx => {
110            result.map_err(|_e| CloudError::operator_auth("sign-in cancelled"))?
111        }
112        _ = server => Err(CloudError::operator_auth("callback server stopped unexpectedly")),
113        () = tokio::time::sleep(Duration::from_secs(CALLBACK_TIMEOUT_SECS)) => {
114            Err(CloudError::operator_auth(format!(
115                "sign-in timed out after {CALLBACK_TIMEOUT_SECS} seconds"
116            )))
117        }
118    }
119}
120
121pub async fn run_operator_login(
122    request: &OperatorLoginRequest,
123    templates: OAuthTemplates,
124) -> CloudResult<OperatorCredentials> {
125    let io = LoginIo {
126        callback_port: CALLBACK_PORT,
127        open_browser: &|url| open::that(url),
128    };
129    run_operator_login_with(request, templates, io).await
130}
131
132pub async fn run_operator_login_with(
133    request: &OperatorLoginRequest,
134    templates: OAuthTemplates,
135    io: LoginIo<'_>,
136) -> CloudResult<OperatorCredentials> {
137    let client = reqwest::Client::builder()
138        .connect_timeout(HTTP_CONNECT_TIMEOUT)
139        .timeout(HTTP_DEFAULT_TIMEOUT)
140        .build()?;
141    let discovery = discover(&client, &request.issuer).await?;
142
143    let verifier = random_token();
144    let state = random_token();
145    let nonce = random_token();
146    let listener = crate::callback_listener::bind_callback_listener(io.callback_port)?;
147    let redirect_uri = format!(
148        "http://127.0.0.1:{}/callback",
149        listener.local_addr()?.port()
150    );
151
152    let (tx, rx) = oneshot::channel();
153    let callback = Arc::new(CallbackState {
154        expected_state: state.clone(),
155        tx: Mutex::new(Some(tx)),
156        success_html: templates.success_html,
157        error_html: templates.error_html,
158    });
159    let app = Router::new()
160        .route("/callback", get(callback_handler))
161        .with_state(callback);
162
163    let url = authorize_url(
164        &discovery,
165        request,
166        AuthorizeParams {
167            redirect_uri: &redirect_uri,
168            state: &state,
169            nonce: &nonce,
170            challenge: &pkce_challenge(&verifier),
171        },
172    );
173    CliService::info(&format!("Opening browser to sign in at {}", request.issuer));
174    if let Err(e) = (io.open_browser)(&url) {
175        CliService::warning(&format!("Could not open browser automatically: {e}"));
176        CliService::key_value("URL", &url);
177    }
178    CliService::info("Waiting for sign-in...");
179
180    let code = await_code(listener, app, rx).await?;
181
182    let id_token = redeem_code(
183        &client,
184        &discovery,
185        request,
186        CodeRedemption {
187            code: &code,
188            redirect_uri: &redirect_uri,
189            verifier: &verifier,
190        },
191    )
192    .await?;
193
194    let claims = identity_claims(&id_token)?;
195    if claims.nonce.as_deref() != Some(nonce.as_str()) {
196        return Err(CloudError::operator_auth("id_token nonce did not match"));
197    }
198    let email = claims
199        .email
200        .map(Email::try_new)
201        .transpose()
202        .map_err(|source| CloudError::OperatorEmail { source })?;
203
204    Ok(OperatorCredentials {
205        issuer: request.issuer.trim_end_matches('/').to_owned(),
206        client_id: request.client_id.clone(),
207        id_token: CloudAuthToken::new(id_token),
208        email,
209        authenticated_at: Utc::now(),
210    })
211}
systemprompt_cloud/operator/login.rs

systemprompt_cloud/operator/
login.rs
