mz_frontegg_mock/handlers/
auth.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use crate::models::*;
11use crate::server::Context;
12use crate::utils::{
13    RefreshTokenTarget, generate_access_token, generate_refresh_token, get_user_groups,
14};
15use axum::{Json, extract::State, http::StatusCode};
16use mz_frontegg_auth::{ApiTokenResponse, ClaimTokenType};
17use serde::{Deserialize, Serialize};
18use std::sync::Arc;
19use std::sync::atomic::Ordering;
20
21pub async fn handle_post_auth_api_token(
22    State(context): State<Arc<Context>>,
23    Json(request): Json<ApiToken>,
24) -> Result<Json<ApiTokenResponse>, StatusCode> {
25    *context.auth_requests.lock().unwrap() += 1;
26
27    if !context.enable_auth.load(Ordering::Relaxed) {
28        return Err(StatusCode::UNAUTHORIZED);
29    }
30
31    let user_api_tokens = context.user_api_tokens.lock().unwrap();
32    let access_token = match user_api_tokens
33        .iter()
34        .find(|(token, _)| token.client_id == request.client_id && token.secret == request.secret)
35    {
36        Some((_, email)) => {
37            let users = context.users.lock().unwrap();
38            let user = users
39                .get(email)
40                .expect("API tokens are only created by logged in valid users.");
41            let groups = get_user_groups(&context, &user.id);
42            generate_access_token(
43                &context,
44                ClaimTokenType::UserApiToken,
45                request.client_id,
46                Some(email.to_owned()),
47                Some(user.id),
48                user.tenant_id,
49                user.roles.clone(),
50                groups,
51                None,
52            )
53        }
54        None => {
55            let tenant_api_tokens = context.tenant_api_tokens.lock().unwrap();
56            match tenant_api_tokens.iter().find(|(token, _)| {
57                token.client_id == request.client_id && token.secret == request.secret
58            }) {
59                Some((_, config)) => generate_access_token(
60                    &context,
61                    ClaimTokenType::TenantApiToken,
62                    request.client_id,
63                    None,
64                    None,
65                    config.tenant_id,
66                    config.roles.clone(),
67                    Vec::new(),
68                    config.metadata.clone(),
69                ),
70                None => return Err(StatusCode::UNAUTHORIZED),
71            }
72        }
73    };
74    let refresh_token = generate_refresh_token(&context, RefreshTokenTarget::ApiToken(request));
75    Ok(Json(ApiTokenResponse {
76        expires: "".to_string(),
77        expires_in: context.expires_in_secs,
78        access_token,
79        refresh_token,
80    }))
81}
82
83pub async fn handle_post_auth_user(
84    State(context): State<Arc<Context>>,
85    Json(request): Json<AuthUserRequest>,
86) -> Result<Json<ApiTokenResponse>, StatusCode> {
87    *context.auth_requests.lock().unwrap() += 1;
88
89    if !context.enable_auth.load(Ordering::Relaxed) {
90        return Err(StatusCode::UNAUTHORIZED);
91    }
92
93    let users = context.users.lock().unwrap();
94    let user = match users.get(&request.email) {
95        Some(user) if request.password == user.password => user.to_owned(),
96        _ => return Err(StatusCode::UNAUTHORIZED),
97    };
98    let groups = get_user_groups(&context, &user.id);
99    let access_token = generate_access_token(
100        &context,
101        ClaimTokenType::UserToken,
102        user.id,
103        Some(request.email.clone()),
104        Some(user.id),
105        user.tenant_id,
106        user.roles.clone(),
107        groups,
108        None,
109    );
110    let refresh_token = generate_refresh_token(&context, RefreshTokenTarget::User(request));
111    Ok(Json(ApiTokenResponse {
112        expires: "".to_string(),
113        expires_in: context.expires_in_secs,
114        access_token,
115        refresh_token,
116    }))
117}
118
119pub async fn handle_post_token_refresh(
120    State(context): State<Arc<Context>>,
121    Json(previous_refresh_token): Json<RefreshTokenRequest>,
122) -> Result<Json<ApiTokenResponse>, StatusCode> {
123    // Always count refresh attempts, even if enable_refresh is false.
124    *context.refreshes.lock().unwrap() += 1;
125
126    if !context.enable_auth.load(Ordering::Relaxed) {
127        return Err(StatusCode::UNAUTHORIZED);
128    }
129
130    let maybe_target = context
131        .refresh_tokens
132        .lock()
133        .unwrap()
134        .remove(&previous_refresh_token.refresh_token);
135    let Some(target) = maybe_target else {
136        return Err(StatusCode::UNAUTHORIZED);
137    };
138
139    match target {
140        RefreshTokenTarget::User(request) => {
141            handle_post_auth_user(State(context), Json(request)).await
142        }
143        RefreshTokenTarget::ApiToken(request) => {
144            handle_post_auth_api_token(State(context), Json(request)).await
145        }
146    }
147}
148
149/// Request body for vendor authentication.
150#[derive(Debug, Clone, Deserialize)]
151#[serde(rename_all = "camelCase")]
152pub struct VendorAuthRequest {
153    pub client_id: String,
154    pub secret: String,
155}
156
157/// Response body for vendor authentication.
158#[derive(Debug, Clone, Serialize)]
159#[serde(rename_all = "camelCase")]
160pub struct VendorAuthResponse {
161    pub token: String,
162    pub expires_in: i64,
163}
164
165/// Handle POST /auth/vendor
166/// Authenticates a vendor SDK client and returns a JWT token.
167pub async fn handle_post_auth_vendor(
168    State(context): State<Arc<Context>>,
169    Json(_request): Json<VendorAuthRequest>,
170) -> Result<Json<VendorAuthResponse>, StatusCode> {
171    // For the mock, we accept any client_id/secret and return a valid token.
172    // In production, this would validate the credentials.
173    let now_secs = context.now.as_secs();
174    let exp_secs = now_secs + context.expires_in_secs;
175    let token = jsonwebtoken::encode(
176        &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256),
177        &serde_json::json!({
178            "sub": "vendor",
179            "iss": context.issuer,
180            "iat": now_secs,
181            "exp": exp_secs,
182        }),
183        &context.encoding_key,
184    )
185    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
186
187    Ok(Json(VendorAuthResponse {
188        token,
189        expires_in: context.expires_in_secs,
190    }))
191}
mz_frontegg_mock/handlers/auth.rs

mz_frontegg_mock/handlers/
auth.rs
