multi tenancy part 1

2025-10-22 22:52:59 +02:00
parent 8e3f09774a
commit e175d28c2c
35 changed files with 1288 additions and 314 deletions
@@ -1,11 +1,15 @@
 use axum::{
    extract::State,
    http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
+    response::{IntoResponse, Response},
    Json,
 };
-use axum_extra::{headers::Cookie, typed_header::TypedHeader};
+use axum_extra::{
+    headers::{authorization::Bearer, Authorization, Cookie},
+    typed_header::TypedHeader,
+};
 use chrono::{Duration as ChronoDuration, Utc};
-use diesel::prelude::*;
+use diesel::{pg::PgConnection, prelude::*};
 use rand::{rngs::OsRng, RngCore};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
@@ -14,8 +18,11 @@ use uuid::Uuid;
 use crate::{
    auth::{password, AuthenticatedUser},
    error::{AppError, AppResult},
-    models::{NewRefreshToken, RefreshToken, User},
-    schema::{refresh_tokens, users::dsl},
+    models::{NewRefreshToken, RefreshToken, Tenant, User, UserMembership},
+    schema::{
+        refresh_tokens, tenants::dsl as tenant_dsl, user_memberships::dsl as memberships_dsl,
+        users::dsl,
+    },
    state::AppState,
 };

@@ -27,6 +34,8 @@ const REFRESH_COOKIE_NAME: &str = "refresh_token";
 pub struct LoginRequest {
    pub username: String,
    pub password: String,
+    #[serde(default)]
+    pub preferred_tenant_slug: Option<String>,
 }

 #[derive(Serialize)]
@@ -36,11 +45,28 @@ pub struct LoginResponse {
    pub expires_in: i64,
 }

+#[derive(Serialize)]
+pub struct TenantSummary {
+    pub tenant_id: Uuid,
+    pub slug: String,
+}
+
+#[derive(Serialize)]
+pub struct TenantSelectionResponse {
+    pub selection_token: String,
+    pub tenants: Vec<TenantSummary>,
+}
+
+#[derive(Deserialize)]
+pub struct TenantSelectionRequest {
+    pub tenant_id: Uuid,
+}
+
 pub async fn login(
    State(state): State<AppState>,
    Json(payload): Json<LoginRequest>,
-) -> AppResult<(HeaderMap, Json<LoginResponse>)> {
-    let mut conn = state.db()?;
+) -> AppResult<Response> {
+    let mut conn = state.db_unscoped()?;

    let user: User = dsl::users
        .filter(dsl::username.eq(&payload.username))
@@ -53,55 +79,67 @@ pub async fn login(
        return Err(AppError::unauthorized());
    }

-    let access_token = state
+    let memberships: Vec<(UserMembership, Tenant)> = memberships_dsl::user_memberships
+        .inner_join(tenant_dsl::tenants)
+        .filter(memberships_dsl::user_id.eq(user.id))
+        .load(&mut conn)?;
+
+    if memberships.is_empty() {
+        return Err(AppError::unauthorized());
+    }
+
+    let preferred_slug = payload
+        .preferred_tenant_slug
+        .as_ref()
+        .map(|slug| slug.trim().to_string())
+        .filter(|slug| !slug.is_empty());
+
+    if let Some(tenant) = preferred_slug.as_ref().and_then(|slug| {
+        memberships
+            .iter()
+            .find(|(_, tenant)| tenant.slug.eq_ignore_ascii_case(slug))
+    }) {
+        return issue_session(&state, &mut conn, &user, tenant.1.tenant_id);
+    }
+
+    if memberships.len() == 1 {
+        let tenant_id = memberships[0].1.tenant_id;
+        return issue_session(&state, &mut conn, &user, tenant_id);
+    }
+
+    let selection_token = state
        .jwt
-        .generate_token(user.id, &user.username, &user.role)
+        .generate_tenant_selector_token(user.id)
        .map_err(AppError::from)?;

-    let now = Utc::now();
-    let refresh_value = generate_refresh_token();
-    let refresh_hash = hash_refresh_token(&refresh_value);
-    let refresh_expires_at = now + ChronoDuration::days(state.config.refresh_token_expiry_days);
+    let tenants = memberships
+        .into_iter()
+        .map(|(_, tenant)| TenantSummary {
+            tenant_id: tenant.tenant_id,
+            slug: tenant.slug,
+        })
+        .collect();

-    let new_refresh = NewRefreshToken {
-        id: Uuid::new_v4(),
-        user_id: user.id,
-        token_hash: refresh_hash,
-        issued_at: now.naive_utc(),
-        expires_at: refresh_expires_at.naive_utc(),
-    };
+    let response = Json(TenantSelectionResponse {
+        selection_token,
+        tenants,
+    })
+    .into_response();

-    diesel::insert_into(refresh_tokens::table)
-        .values(&new_refresh)
-        .execute(&mut conn)?;
-
-    let mut headers = HeaderMap::new();
-    headers.insert(
-        SET_COOKIE,
-        build_refresh_cookie(&state, &refresh_value, refresh_expires_at),
-    );
-
-    Ok((
-        headers,
-        Json(LoginResponse {
-            access_token,
-            token_type: "Bearer".to_string(),
-            expires_in: state.config.jwt_expiry_minutes * 60,
-        }),
-    ))
+    Ok(response)
 }

 pub async fn refresh(
    State(state): State<AppState>,
    jar: Option<TypedHeader<Cookie>>,
-) -> AppResult<(HeaderMap, Json<LoginResponse>)> {
+) -> AppResult<Response> {
    let cookies = jar.ok_or_else(AppError::unauthorized)?;
    let refresh_value = cookies
        .get(REFRESH_COOKIE_NAME)
        .ok_or_else(AppError::unauthorized)?;

    let hashed = hash_refresh_token(refresh_value);
-    let mut conn = state.db()?;
+    let mut conn = state.db_unscoped()?;
    let now = Utc::now();
    let now_naive = now.naive_utc();

@@ -128,41 +166,39 @@ pub async fn refresh(
        .first(&mut conn)
        .map_err(AppError::from)?;

-    let access_token = state
+    issue_session(&state, &mut conn, &user, token.tenant_id)
+}
+
+pub async fn select_tenant(
+    State(state): State<AppState>,
+    TypedHeader(Authorization(bearer)): TypedHeader<Authorization<Bearer>>,
+    Json(payload): Json<TenantSelectionRequest>,
+) -> AppResult<Response> {
+    let claims = state
        .jwt
-        .generate_token(user.id, &user.username, &user.role)
+        .verify_tenant_selector_token(bearer.token())
+        .map_err(|_| AppError::unauthorized())?;
+
+    let mut conn = state.db_unscoped()?;
+
+    let membership_exists = memberships_dsl::user_memberships
+        .filter(memberships_dsl::user_id.eq(claims.sub))
+        .filter(memberships_dsl::tenant_id.eq(payload.tenant_id))
+        .inner_join(tenant_dsl::tenants)
+        .select(memberships_dsl::id)
+        .first::<Uuid>(&mut conn)
+        .optional()?;
+
+    if membership_exists.is_none() {
+        return Err(AppError::unauthorized());
+    }
+
+    let user: User = dsl::users
+        .find(claims.sub)
+        .first(&mut conn)
        .map_err(AppError::from)?;

-    let new_refresh_value = generate_refresh_token();
-    let new_refresh_hash = hash_refresh_token(&new_refresh_value);
-    let new_refresh_expires = now + ChronoDuration::days(state.config.refresh_token_expiry_days);
-
-    let new_refresh = NewRefreshToken {
-        id: Uuid::new_v4(),
-        user_id: user.id,
-        token_hash: new_refresh_hash,
-        issued_at: now_naive,
-        expires_at: new_refresh_expires.naive_utc(),
-    };
-
-    diesel::insert_into(refresh_tokens::table)
-        .values(&new_refresh)
-        .execute(&mut conn)?;
-
-    let mut headers = HeaderMap::new();
-    headers.insert(
-        SET_COOKIE,
-        build_refresh_cookie(&state, &new_refresh_value, new_refresh_expires),
-    );
-
-    Ok((
-        headers,
-        Json(LoginResponse {
-            access_token,
-            token_type: "Bearer".to_string(),
-            expires_in: state.config.jwt_expiry_minutes * 60,
-        }),
-    ))
+    issue_session(&state, &mut conn, &user, payload.tenant_id)
 }

 pub async fn logout(
@@ -170,7 +206,7 @@ pub async fn logout(
    user: AuthenticatedUser,
    jar: Option<TypedHeader<Cookie>>,
 ) -> AppResult<(HeaderMap, StatusCode)> {
-    let mut conn = state.db()?;
+    let mut conn = state.db_unscoped()?;
    let now = Utc::now().naive_utc();
    let mut rows_affected = 0;

@@ -214,6 +250,50 @@ pub async fn me(user: AuthenticatedUser) -> Json<AuthenticatedUser> {
    Json(user)
 }

+fn issue_session(
+    state: &AppState,
+    conn: &mut PgConnection,
+    user: &User,
+    tenant_id: Uuid,
+) -> AppResult<Response> {
+    let now = Utc::now();
+    let access_token = state
+        .jwt
+        .generate_token(user.id, tenant_id, &user.username)
+        .map_err(AppError::from)?;
+
+    let refresh_value = generate_refresh_token();
+    let refresh_hash = hash_refresh_token(&refresh_value);
+    let refresh_expires_at = now + ChronoDuration::days(state.config.refresh_token_expiry_days);
+
+    let new_refresh = NewRefreshToken {
+        id: Uuid::new_v4(),
+        user_id: user.id,
+        token_hash: refresh_hash,
+        issued_at: now.naive_utc(),
+        expires_at: refresh_expires_at.naive_utc(),
+        tenant_id,
+    };
+
+    diesel::insert_into(refresh_tokens::table)
+        .values(&new_refresh)
+        .execute(conn)?;
+
+    let mut response = Json(LoginResponse {
+        access_token,
+        token_type: "Bearer".to_string(),
+        expires_in: state.config.jwt_expiry_minutes * 60,
+    })
+    .into_response();
+
+    response.headers_mut().insert(
+        SET_COOKIE,
+        build_refresh_cookie(state, &refresh_value, refresh_expires_at),
+    );
+
+    Ok(response)
+}
+
 fn hash_refresh_token(token: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());