use std::sync::Arc;
use std::time::Duration;
use axum::Router;
use axum::extract::{Query, State};
use axum::response::Html;
use axum::routing::get;
use base64::prelude::*;
use chrono::Utc;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use systemprompt_identifiers::{ClientId, CloudAuthToken, Email};
use systemprompt_logging::CliService;
use systemprompt_models::net::{HTTP_CONNECT_TIMEOUT, HTTP_DEFAULT_TIMEOUT};
use tokio::sync::{Mutex, oneshot};
use uuid::Uuid;
use super::endpoints::{AuthorizeParams, CodeRedemption, authorize_url, discover, redeem_code};
use super::{OperatorCredentials, identity_claims};
use crate::OAuthTemplates;
use crate::constants::oauth::{CALLBACK_PORT, CALLBACK_TIMEOUT_SECS};
use crate::error::{CloudError, CloudResult};
#[derive(Debug, Clone)]
pub struct OperatorLoginRequest {
pub issuer: String,
pub client_id: ClientId,
pub scopes: Vec<String>,
}
pub struct LoginIo<'a> {
pub callback_port: u16,
pub open_browser: &'a (dyn Fn(&str) -> std::io::Result<()> + Sync),
}
impl std::fmt::Debug for LoginIo<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LoginIo")
.field("callback_port", &self.callback_port)
.finish_non_exhaustive()
}
}
#[derive(Deserialize)]
struct CallbackParams {
code: Option<String>,
state: Option<String>,
error: Option<String>,
error_description: Option<String>,
}
struct CallbackState {
expected_state: String,
tx: Mutex<Option<oneshot::Sender<CloudResult<String>>>>,
success_html: &'static str,
error_html: &'static str,
}
fn random_token() -> String {
format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple())
}
fn pkce_challenge(verifier: &str) -> String {
BASE64_URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()))
}
async fn callback_handler(
State(state): State<Arc<CallbackState>>,
Query(params): Query<CallbackParams>,
) -> Html<&'static str> {
let result = if let Some(error) = params.error {
Err(CloudError::operator_auth(format!(
"{error}: {}",
params.error_description.unwrap_or_default()
)))
} else if params.state.as_deref() != Some(state.expected_state.as_str()) {
Err(CloudError::operator_auth("callback state did not match"))
} else {
params
.code
.ok_or_else(|| CloudError::operator_auth("callback carried no authorization code"))
};
let is_success = result.is_ok();
let sender = state.tx.lock().await.take();
if let Some(sender) = sender
&& sender.send(result).is_err()
{
tracing::warn!("Operator login receiver dropped before the callback arrived");
}
Html(if is_success {
state.success_html
} else {
state.error_html
})
}
async fn await_code(
listener: tokio::net::TcpListener,
app: Router,
rx: oneshot::Receiver<CloudResult<String>>,
) -> CloudResult<String> {
let server = axum::serve(listener, app);
tokio::select! {
result = rx => {
result.map_err(|_e| CloudError::operator_auth("sign-in cancelled"))?
}
_ = server => Err(CloudError::operator_auth("callback server stopped unexpectedly")),
() = tokio::time::sleep(Duration::from_secs(CALLBACK_TIMEOUT_SECS)) => {
Err(CloudError::operator_auth(format!(
"sign-in timed out after {CALLBACK_TIMEOUT_SECS} seconds"
)))
}
}
}
pub async fn run_operator_login(
request: &OperatorLoginRequest,
templates: OAuthTemplates,
) -> CloudResult<OperatorCredentials> {
let io = LoginIo {
callback_port: CALLBACK_PORT,
open_browser: &|url| open::that(url),
};
run_operator_login_with(request, templates, io).await
}
pub async fn run_operator_login_with(
request: &OperatorLoginRequest,
templates: OAuthTemplates,
io: LoginIo<'_>,
) -> CloudResult<OperatorCredentials> {
let client = reqwest::Client::builder()
.connect_timeout(HTTP_CONNECT_TIMEOUT)
.timeout(HTTP_DEFAULT_TIMEOUT)
.build()?;
let discovery = discover(&client, &request.issuer).await?;
let verifier = random_token();
let state = random_token();
let nonce = random_token();
let listener = crate::callback_listener::bind_callback_listener(io.callback_port)?;
let redirect_uri = format!(
"http://127.0.0.1:{}/callback",
listener.local_addr()?.port()
);
let (tx, rx) = oneshot::channel();
let callback = Arc::new(CallbackState {
expected_state: state.clone(),
tx: Mutex::new(Some(tx)),
success_html: templates.success_html,
error_html: templates.error_html,
});
let app = Router::new()
.route("/callback", get(callback_handler))
.with_state(callback);
let url = authorize_url(
&discovery,
request,
AuthorizeParams {
redirect_uri: &redirect_uri,
state: &state,
nonce: &nonce,
challenge: &pkce_challenge(&verifier),
},
);
CliService::info(&format!("Opening browser to sign in at {}", request.issuer));
if let Err(e) = (io.open_browser)(&url) {
CliService::warning(&format!("Could not open browser automatically: {e}"));
CliService::key_value("URL", &url);
}
CliService::info("Waiting for sign-in...");
let code = await_code(listener, app, rx).await?;
let id_token = redeem_code(
&client,
&discovery,
request,
CodeRedemption {
code: &code,
redirect_uri: &redirect_uri,
verifier: &verifier,
},
)
.await?;
let claims = identity_claims(&id_token)?;
if claims.nonce.as_deref() != Some(nonce.as_str()) {
return Err(CloudError::operator_auth("id_token nonce did not match"));
}
let email = claims
.email
.map(Email::try_new)
.transpose()
.map_err(|source| CloudError::OperatorEmail { source })?;
Ok(OperatorCredentials {
issuer: request.issuer.trim_end_matches('/').to_owned(),
client_id: request.client_id.clone(),
id_token: CloudAuthToken::new(id_token),
email,
authenticated_at: Utc::now(),
})
}