multi tenancy part 1

backend/src/auth/jwt.rs

@@ -15,6 +15,8 @@ pub struct JwtService {
     expiry: Duration,
     download_audience: String,
     download_expiry: Duration,
+    selector_audience: String,
+    selector_expiry: Duration,
 }
 
 impl JwtService {
@@ -27,16 +29,18 @@ impl JwtService {
             expiry: Duration::minutes(config.jwt_expiry_minutes),
             download_audience: config.download_token_audience.clone(),
             download_expiry: Duration::minutes(config.download_token_expiry_minutes),
+            selector_audience: format!("{}:tenant-selector", config.jwt_audience),
+            selector_expiry: Duration::minutes(15),
         })
     }
 
-    pub fn generate_token(&self, user_id: Uuid, username: &str, role: &str) -> Result<String> {
+    pub fn generate_token(&self, user_id: Uuid, tenant_id: Uuid, username: &str) -> Result<String> {
         let now = Utc::now();
         let exp = now + self.expiry;
         let claims = Claims {
             sub: user_id,
+            tenant_id,
             username: username.to_owned(),
-            role: role.to_owned(),
             iss: self.issuer.clone(),
             aud: self.audience.clone(),
             iat: now.timestamp() as usize,
@@ -54,12 +58,18 @@ impl JwtService {
         Ok(data.claims)
     }
 
-    pub fn generate_download_token(&self, document_id: Uuid, user_id: Uuid) -> Result<String> {
+    pub fn generate_download_token(
+        &self,
+        document_id: Uuid,
+        user_id: Uuid,
+        tenant_id: Uuid,
+    ) -> Result<String> {
         let now = Utc::now();
         let exp = now + self.download_expiry;
         let claims = DownloadClaims {
             doc_id: document_id,
             user_id,
+            tenant_id,
             iss: self.issuer.clone(),
             aud: self.download_audience.clone(),
             iat: now.timestamp() as usize,
@@ -76,13 +86,35 @@ impl JwtService {
         let data = decode::<DownloadClaims>(token, &self.decoding, &validation)?;
         Ok(data.claims)
     }
+
+    pub fn generate_tenant_selector_token(&self, user_id: Uuid) -> Result<String> {
+        let now = Utc::now();
+        let exp = now + self.selector_expiry;
+        let claims = TenantSelectionClaims {
+            sub: user_id,
+            iss: self.issuer.clone(),
+            aud: self.selector_audience.clone(),
+            iat: now.timestamp() as usize,
+            exp: exp.timestamp() as usize,
+        };
+
+        Ok(encode(&Header::default(), &claims, &self.encoding)?)
+    }
+
+    pub fn verify_tenant_selector_token(&self, token: &str) -> Result<TenantSelectionClaims> {
+        let mut validation = Validation::default();
+        validation.set_audience(&[self.selector_audience.clone()]);
+        validation.set_issuer(&[self.issuer.clone()]);
+        let data = decode::<TenantSelectionClaims>(token, &self.decoding, &validation)?;
+        Ok(data.claims)
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Claims {
     pub sub: Uuid,
+    pub tenant_id: Uuid,
     pub username: String,
-    pub role: String,
     pub iss: String,
     pub aud: String,
     pub iat: usize,
@@ -93,6 +125,16 @@ pub struct Claims {
 pub struct DownloadClaims {
     pub doc_id: Uuid,
     pub user_id: Uuid,
+    pub tenant_id: Uuid,
+    pub iss: String,
+    pub aud: String,
+    pub iat: usize,
+    pub exp: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TenantSelectionClaims {
+    pub sub: Uuid,
     pub iss: String,
     pub aud: String,
     pub iat: usize,

backend/src/auth/mod.rs

@@ -6,13 +6,17 @@ use axum_extra::headers::{authorization::Bearer, Authorization};
 use axum_extra::TypedHeader;
 use serde::{Deserialize, Serialize};
 
-use crate::{error::AppError, state::AppState};
+use crate::{
+    error::AppError,
+    state::{AppState, PgPooledConnection},
+};
+use uuid::Uuid;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct AuthenticatedUser {
     pub user_id: uuid::Uuid,
     pub username: String,
-    pub role: String,
+    pub tenant_id: uuid::Uuid,
 }
 
 #[async_trait]
@@ -23,6 +27,10 @@ impl FromRequestParts<AppState> for AuthenticatedUser {
         parts: &mut Parts,
         state: &AppState,
     ) -> Result<Self, Self::Rejection> {
+        if let Some(user) = parts.extensions.get::<AuthenticatedUser>() {
+            return Ok(user.clone());
+        }
+
         let TypedHeader(Authorization(bearer)) =
             TypedHeader::<Authorization<Bearer>>::from_request_parts(parts, state)
                 .await
@@ -33,10 +41,48 @@ impl FromRequestParts<AppState> for AuthenticatedUser {
             .verify_token(bearer.token())
             .map_err(|_| AppError::unauthorized())?;
 
-        Ok(AuthenticatedUser {
+        let user = AuthenticatedUser {
             user_id: claims.sub,
             username: claims.username,
-            role: claims.role,
+            tenant_id: claims.tenant_id,
+        };
+
+        parts.extensions.insert(user.clone());
+
+        Ok(user)
+    }
+}
+
+pub struct TenantScopedConn {
+    pub conn: PgPooledConnection,
+    pub tenant_id: Uuid,
+    pub user_id: Uuid,
+    pub user: AuthenticatedUser,
+}
+
+impl TenantScopedConn {
+    pub fn conn(&mut self) -> &mut PgPooledConnection {
+        &mut self.conn
+    }
+}
+
+#[async_trait]
+impl FromRequestParts<AppState> for TenantScopedConn {
+    type Rejection = AppError;
+
+    async fn from_request_parts(
+        parts: &mut Parts,
+        state: &AppState,
+    ) -> Result<Self, Self::Rejection> {
+        let user = AuthenticatedUser::from_request_parts(parts, state).await?;
+        let tenant_id = user.tenant_id;
+        let conn = state.db_for_tenant(tenant_id)?;
+
+        Ok(Self {
+            conn,
+            tenant_id,
+            user_id: user.user_id,
+            user,
         })
     }
 }