systemprompt_api/routes/oauth/endpoints/
callback.rs

1//! OAuth callback endpoint for the server's own browser client.
2//!
3//! Exchanges the returned authorization code for tokens, establishes an
4//! authenticated session, sets the access-token cookie, and redirects to the
5//! origin-validated `return_to` recovered from the consumed state binding.
6//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10use axum::extract::{Query, State};
11use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
12use axum::response::{IntoResponse, Redirect, Response};
13use serde::Deserialize;
14use std::str::FromStr;
15use std::sync::Arc;
16use systemprompt_identifiers::{
17    AuthorizationCode, ClientId, RefreshTokenId, SessionSource, UserId,
18};
19use systemprompt_models::Config;
20use systemprompt_models::auth::{AuthenticatedUser, Permission, parse_permissions};
21
22use crate::routes::oauth::extractors::OAuthRepo;
23use systemprompt_oauth::OAuthState;
24use systemprompt_oauth::repository::{OAuthRepository, RefreshTokenParams};
25
26#[derive(Debug, Deserialize)]
27pub struct CallbackQuery {
28    pub code: String,
29    pub state: Option<String>,
30}
31
32pub async fn handle_callback(
33    Query(params): Query<CallbackQuery>,
34    State(state): State<OAuthState>,
35    OAuthRepo(repo): OAuthRepo,
36    headers: HeaderMap,
37) -> impl IntoResponse {
38    let config = match Config::get() {
39        Ok(c) => c,
40        Err(e) => {
41            return (
42                StatusCode::INTERNAL_SERVER_ERROR,
43                format!("Failed to load config: {e}"),
44            )
45                .into_response();
46        },
47    };
48
49    let server_base_url = &config.api_external_url;
50    let redirect_uri = format!("{server_base_url}/api/v1/core/oauth/callback");
51
52    let browser_client = match find_browser_client(&repo, &redirect_uri).await {
53        Ok(client) => client,
54        Err(e) => {
55            return (
56                StatusCode::INTERNAL_SERVER_ERROR,
57                format!("Failed to find OAuth client: {e}"),
58            )
59                .into_response();
60        },
61    };
62
63    let code = AuthorizationCode::new(&params.code);
64    let client_id = ClientId::new(&browser_client.client_id);
65    let token_response = match exchange_code_for_token(
66        &repo,
67        CodeExchangeParams {
68            code: &code,
69            client_id: &client_id,
70            redirect_uri: &redirect_uri,
71            headers: &headers,
72        },
73        &state,
74    )
75    .await
76    {
77        Ok(response) => response,
78        Err(e) => {
79            return (
80                StatusCode::UNAUTHORIZED,
81                format!("Failed to exchange code for token: {e}"),
82            )
83                .into_response();
84        },
85    };
86
87    let redirect_destination =
88        match resolve_redirect_destination(&repo, params.state.as_deref()).await {
89            Ok(destination) => destination,
90            Err(response) => return response,
91        };
92
93    session_cookie_redirect(&token_response.access_token, &redirect_destination)
94}
95
96async fn resolve_redirect_destination(
97    repo: &OAuthRepository,
98    state_token: Option<&str>,
99) -> Result<String, Response> {
100    let Some(state_token) = state_token.filter(|s| !s.is_empty()) else {
101        return Err((StatusCode::BAD_REQUEST, "Missing state parameter").into_response());
102    };
103    match repo.consume_state_binding(state_token).await {
104        Ok(Some(binding)) => Ok(binding.return_to),
105        Ok(None) => {
106            tracing::warn!("state binding missing, expired, or already consumed");
107            Err((StatusCode::BAD_REQUEST, "Invalid state parameter").into_response())
108        },
109        Err(e) => {
110            tracing::error!(error = %e, "state binding lookup failed");
111            Err((
112                StatusCode::INTERNAL_SERVER_ERROR,
113                "Failed to validate state",
114            )
115                .into_response())
116        },
117    }
118}
119
120fn session_cookie_redirect(access_token: &str, destination: &str) -> Response {
121    let cookie = format!(
122        "access_token={access_token}; Path=/; HttpOnly; Secure; SameSite=Strict; Max-Age={}",
123        systemprompt_oauth::constants::token::COOKIE_MAX_AGE_SECONDS
124    );
125
126    let mut response = Redirect::to(destination).into_response();
127    if let Ok(cookie_value) = HeaderValue::from_str(&cookie) {
128        response
129            .headers_mut()
130            .insert(header::SET_COOKIE, cookie_value);
131    }
132
133    response
134}
135
136async fn find_browser_client(
137    repo: &OAuthRepository,
138    redirect_uri: &str,
139) -> anyhow::Result<BrowserClient> {
140    let client = repo
141        .find_client_by_redirect_uri_with_scope(redirect_uri, &["admin", "user"])
142        .await?
143        .ok_or_else(|| anyhow::anyhow!("No suitable browser client found"))?;
144
145    Ok(BrowserClient {
146        client_id: client.client_id.to_string(),
147    })
148}
149
150struct CodeExchangeParams<'a> {
151    code: &'a AuthorizationCode,
152    client_id: &'a ClientId,
153    redirect_uri: &'a str,
154    headers: &'a HeaderMap,
155}
156
157async fn exchange_code_for_token(
158    repo: &OAuthRepository,
159    params: CodeExchangeParams<'_>,
160    state: &OAuthState,
161) -> anyhow::Result<TokenResponse> {
162    use systemprompt_oauth::services::{
163        JwtConfig, JwtSigningParams, generate_access_token_jti, generate_jwt, generate_secure_token,
164    };
165
166    let validation_result = repo
167        .validate_authorization_code(
168            params.code,
169            params.client_id,
170            Some(params.redirect_uri),
171            None,
172        )
173        .await?;
174
175    let user = load_authenticated_user(&validation_result.user_id, state.user_provider()).await?;
176
177    let permissions = parse_permissions(&validation_result.scope)?;
178
179    let mut session_service = systemprompt_oauth::services::SessionCreationService::new(
180        Arc::clone(state.analytics_provider()),
181        Arc::clone(state.user_provider()),
182    );
183    if let Some(publisher) = state.event_publisher() {
184        session_service = session_service.with_event_publisher(Arc::clone(publisher));
185    }
186    let session_id = session_service
187        .create_authenticated_session(
188            &validation_result.user_id,
189            params.headers,
190            SessionSource::Oauth,
191        )
192        .await?;
193
194    let access_token_jti = generate_access_token_jti();
195    let global_config = Config::get()?;
196    let config = JwtConfig {
197        permissions: permissions.clone(),
198        audience: global_config.jwt_audiences.clone(),
199        ..Default::default()
200    };
201    let signing = JwtSigningParams {
202        issuer: &global_config.jwt_issuer,
203    };
204    let access_token = generate_jwt(&user, config, access_token_jti, &session_id, &signing)?;
205
206    let refresh_token_value = generate_secure_token("rt");
207    let refresh_token_id = RefreshTokenId::new(&refresh_token_value);
208    let refresh_expires_at = chrono::Utc::now().timestamp()
209        + (systemprompt_oauth::constants::token::SECONDS_PER_DAY
210            * systemprompt_oauth::constants::token::REFRESH_TOKEN_EXPIRY_DAYS);
211
212    let refresh_params = RefreshTokenParams::builder(
213        &refresh_token_id,
214        params.client_id,
215        &validation_result.user_id,
216        &validation_result.scope,
217        refresh_expires_at,
218    )
219    .build();
220    repo.store_refresh_token(refresh_params).await?;
221
222    if let Err(e) = repo
223        .link_auth_code_to_refresh_token(params.code, refresh_token_id.as_str())
224        .await
225    {
226        tracing::warn!(error = %e, "Failed to link auth code to refresh token");
227    }
228
229    Ok(TokenResponse { access_token })
230}
231
232async fn load_authenticated_user(
233    user_id: &UserId,
234    user_provider: &Arc<dyn systemprompt_traits::UserProvider>,
235) -> anyhow::Result<AuthenticatedUser> {
236    let user = user_provider
237        .find_by_id(user_id)
238        .await
239        .map_err(|e| anyhow::anyhow!("{}", e))?
240        .ok_or_else(|| anyhow::anyhow!("User not found: {user_id}"))?;
241
242    let permissions: Vec<Permission> = user
243        .roles
244        .iter()
245        .filter_map(|s| {
246            Permission::from_str(s)
247                .map_err(|e| {
248                    tracing::warn!(
249                        user_id = %user.id,
250                        role = %s,
251                        error = %e,
252                        "Invalid role in user record"
253                    );
254                    e
255                })
256                .ok()
257        })
258        .collect();
259
260    let user_uuid = uuid::Uuid::parse_str(user.id.as_str())
261        .map_err(|_e| anyhow::anyhow!("Invalid user UUID: {}", user.id))?;
262
263    Ok(AuthenticatedUser::new_with_roles(
264        user_uuid,
265        user.name,
266        user.email,
267        permissions,
268        user.roles,
269    ))
270}
271
272#[derive(Debug)]
273struct BrowserClient {
274    client_id: String,
275}
276
277#[derive(Debug, serde::Deserialize)]
278struct TokenResponse {
279    access_token: String,
280}
systemprompt_api/routes/oauth/endpoints/callback.rs

systemprompt_api/routes/oauth/endpoints/
callback.rs
