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
36#[derive(Deserialize)]
37struct CallbackParams {
38    code: Option<String>,
39    state: Option<String>,
40    error: Option<String>,
41    error_description: Option<String>,
42}
43
44struct CallbackState {
45    expected_state: String,
46    tx: Mutex<Option<oneshot::Sender<CloudResult<String>>>>,
47    success_html: &'static str,
48    error_html: &'static str,
49}
50
51fn random_token() -> String {
52    format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple())
53}
54
55fn pkce_challenge(verifier: &str) -> String {
56    BASE64_URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()))
57}
58
59async fn callback_handler(
60    State(state): State<Arc<CallbackState>>,
61    Query(params): Query<CallbackParams>,
62) -> Html<&'static str> {
63    let result = if let Some(error) = params.error {
64        Err(CloudError::operator_auth(format!(
65            "{error}: {}",
66            params.error_description.unwrap_or_default()
67        )))
68    } else if params.state.as_deref() != Some(state.expected_state.as_str()) {
69        Err(CloudError::operator_auth("callback state did not match"))
70    } else {
71        params
72            .code
73            .ok_or_else(|| CloudError::operator_auth("callback carried no authorization code"))
74    };
75    let is_success = result.is_ok();
76    let sender = state.tx.lock().await.take();
77    if let Some(sender) = sender
78        && sender.send(result).is_err()
79    {
80        tracing::warn!("Operator login receiver dropped before the callback arrived");
81    }
82    Html(if is_success {
83        state.success_html
84    } else {
85        state.error_html
86    })
87}
88
89async fn await_code(
90    listener: tokio::net::TcpListener,
91    app: Router,
92    rx: oneshot::Receiver<CloudResult<String>>,
93) -> CloudResult<String> {
94    let server = axum::serve(listener, app);
95    tokio::select! {
96        result = rx => {
97            result.map_err(|_e| CloudError::operator_auth("sign-in cancelled"))?
98        }
99        _ = server => Err(CloudError::operator_auth("callback server stopped unexpectedly")),
100        () = tokio::time::sleep(Duration::from_secs(CALLBACK_TIMEOUT_SECS)) => {
101            Err(CloudError::operator_auth(format!(
102                "sign-in timed out after {CALLBACK_TIMEOUT_SECS} seconds"
103            )))
104        }
105    }
106}
107
108pub async fn run_operator_login(
109    request: &OperatorLoginRequest,
110    templates: OAuthTemplates,
111) -> CloudResult<OperatorCredentials> {
112    let client = reqwest::Client::builder()
113        .connect_timeout(HTTP_CONNECT_TIMEOUT)
114        .timeout(HTTP_DEFAULT_TIMEOUT)
115        .build()?;
116    let discovery = discover(&client, &request.issuer).await?;
117
118    let verifier = random_token();
119    let state = random_token();
120    let nonce = random_token();
121    let redirect_uri = format!("http://127.0.0.1:{CALLBACK_PORT}/callback");
122
123    let (tx, rx) = oneshot::channel();
124    let callback = Arc::new(CallbackState {
125        expected_state: state.clone(),
126        tx: Mutex::new(Some(tx)),
127        success_html: templates.success_html,
128        error_html: templates.error_html,
129    });
130    let app = Router::new()
131        .route("/callback", get(callback_handler))
132        .with_state(callback);
133    let listener = crate::callback_listener::bind_callback_listener(CALLBACK_PORT)?;
134
135    let url = authorize_url(
136        &discovery,
137        request,
138        AuthorizeParams {
139            redirect_uri: &redirect_uri,
140            state: &state,
141            nonce: &nonce,
142            challenge: &pkce_challenge(&verifier),
143        },
144    );
145    CliService::info(&format!("Opening browser to sign in at {}", request.issuer));
146    if let Err(e) = open::that(&url) {
147        CliService::warning(&format!("Could not open browser automatically: {e}"));
148        CliService::key_value("URL", &url);
149    }
150    CliService::info("Waiting for sign-in...");
151
152    let code = await_code(listener, app, rx).await?;
153
154    let id_token = redeem_code(
155        &client,
156        &discovery,
157        request,
158        CodeRedemption {
159            code: &code,
160            redirect_uri: &redirect_uri,
161            verifier: &verifier,
162        },
163    )
164    .await?;
165
166    let claims = identity_claims(&id_token)?;
167    if claims.nonce.as_deref() != Some(nonce.as_str()) {
168        return Err(CloudError::operator_auth("id_token nonce did not match"));
169    }
170    let email = claims
171        .email
172        .map(Email::try_new)
173        .transpose()
174        .map_err(|source| CloudError::OperatorEmail { source })?;
175
176    Ok(OperatorCredentials {
177        issuer: request.issuer.trim_end_matches('/').to_owned(),
178        client_id: request.client_id.clone(),
179        id_token: CloudAuthToken::new(id_token),
180        email,
181        authenticated_at: Utc::now(),
182    })
183}
systemprompt_cloud/operator/login.rs

systemprompt_cloud/operator/
login.rs
