rullst_core/
multitenant.rs

1use std::cell::RefCell;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4#[non_exhaustive]
5/// Strategy used to extract the active tenant ID from an incoming request.
6pub enum TenantStrategy {
7    /// Select a tenant from the request host subdomain (e.g. `tenant.example.com`).
8    /// The selection is accepted only when authenticated membership allows it.
9    Subdomain,
10    /// Select a tenant from a custom HTTP header. The header is an untrusted
11    /// hint and is accepted only when authenticated membership allows it.
12    Header,
13    /// Extract the tenant ID from query parameters.
14    ///
15    /// # Security Warning
16    /// Using query parameters (`?tenant_id=...`) can cause tenant IDs to leak in:
17    /// - Server access logs (clear-text)
18    /// - Browser history
19    /// - `Referer` headers sent to third-party assets
20    /// - CDN/proxy cache headers
21    ///
22    /// For sensitive or regulated tenant IDs, prefer `TenantStrategy::Header` or `TenantStrategy::Subdomain`.
23    /// Every strategy still requires a trusted [`crate::security::TenantMembership`]
24    /// request extension from authentication middleware.
25    Parameter,
26}
27
28#[derive(Debug, Clone)]
29#[non_exhaustive]
30/// Configuration settings for multitenancy extraction.
31pub struct TenantConfig {
32    /// The extraction strategy to be used.
33    pub strategy: TenantStrategy,
34    /// The name of the custom HTTP header (used only with `TenantStrategy::Header`).
35    pub header_name: String,
36    /// The name of the query parameter (used only with `TenantStrategy::Parameter`).
37    pub parameter_name: String,
38    /// Fallback tenant ID to use when subdomain/header resolution fails or is absent.
39    pub domain_fallback: Option<String>,
40}
41
42impl TenantConfig {
43    /// Creates a new TenantConfig with default values:
44    /// - Header Name: X-Tenant-ID
45    /// - Parameter Name: tenant_id
46    /// - Domain Fallback: None
47    pub fn new(strategy: TenantStrategy) -> Self {
48        Self {
49            strategy,
50            header_name: "X-Tenant-ID".to_string(),
51            parameter_name: "tenant_id".to_string(),
52            domain_fallback: None,
53        }
54    }
55
56    /// Set a custom HTTP header name to extract the tenant ID from.
57    pub fn with_header_name<S: Into<String>>(mut self, name: S) -> Self {
58        self.header_name = name.into();
59        self
60    }
61
62    /// Set a custom Query Parameter name to extract the tenant ID from.
63    pub fn with_parameter_name<S: Into<String>>(mut self, name: S) -> Self {
64        self.parameter_name = name.into();
65        self
66    }
67
68    /// Set a fallback tenant ID when domain extraction fails in Subdomain strategy.
69    pub fn with_domain_fallback<S: Into<String>>(mut self, fallback: S) -> Self {
70        self.domain_fallback = Some(fallback.into());
71        self
72    }
73}
74
75tokio::task_local! {
76    /// Request-scoped, thread-safe context for storing the active tenant ID.
77    pub static TENANT_CONTEXT: RefCell<Option<String>>;
78}
79
80/// Retrieve the active request's tenant ID if configured.
81pub fn current_tenant_id() -> Option<String> {
82    TENANT_CONTEXT
83        .try_with(|ctx| ctx.borrow().clone())
84        .unwrap_or(None)
85}
86
87/// Dynamically sets or replaces the tenant ID for the duration of the current task context.
88pub fn set_tenant_id(tenant_id: Option<String>) {
89    let _ = TENANT_CONTEXT.try_with(|ctx| {
90        *ctx.borrow_mut() = tenant_id;
91    });
92}
93
94/// Helper function to extract subdomain from Host header
95fn extract_subdomain(host: &str) -> Option<String> {
96    let host_only = host.split(':').next()?;
97    if host_only.parse::<std::net::IpAddr>().is_ok() {
98        return None;
99    }
100    let parts: Vec<&str> = host_only.split('.').collect();
101    if parts.len() >= 3 {
102        // e.g. tenant1.example.com -> tenant1
103        Some(parts[0].to_string())
104    } else {
105        None
106    }
107}
108
109/// The declarative custom Tower Layer for tenant identification
110#[derive(Clone, Debug)]
111pub struct TenantLayer {
112    config: TenantConfig,
113}
114
115impl TenantLayer {
116    /// Creates a new `TenantLayer` with the specified `TenantConfig`.
117    pub fn new(config: TenantConfig) -> Self {
118        Self { config }
119    }
120}
121
122impl<S> tower_layer::Layer<S> for TenantLayer {
123    type Service = TenantService<S>;
124
125    fn layer(&self, inner: S) -> Self::Service {
126        TenantService {
127            inner,
128            config: self.config.clone(),
129        }
130    }
131}
132
133/// The declarative custom Tower Service for tenant identification
134#[derive(Clone, Debug)]
135pub struct TenantService<S> {
136    inner: S,
137    config: TenantConfig,
138}
139
140impl<S, ReqBody, ResBody> tower_service::Service<axum::http::Request<ReqBody>> for TenantService<S>
141where
142    S: tower_service::Service<
143            axum::http::Request<ReqBody>,
144            Response = axum::http::Response<ResBody>,
145        > + Clone
146        + Send
147        + 'static,
148    S::Future: Send + 'static,
149    ReqBody: Send + 'static,
150    ResBody: Default + Send + 'static,
151{
152    type Response = axum::http::Response<ResBody>;
153    type Error = S::Error;
154    type Future = std::pin::Pin<
155        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
156    >;
157
158    fn poll_ready(
159        &mut self,
160        cx: &mut std::task::Context<'_>,
161    ) -> std::task::Poll<Result<(), Self::Error>> {
162        self.inner.poll_ready(cx)
163    }
164
165    fn call(&mut self, mut req: axum::http::Request<ReqBody>) -> Self::Future {
166        let config = self.config.clone();
167        // `poll_ready` readied `self.inner`, so that instance must handle the
168        // request: services such as `ConcurrencyLimit` reserve capacity there
169        // and reject a call on a fresh clone. Leave the clone for the next call.
170        let clone = self.inner.clone();
171        let mut inner = std::mem::replace(&mut self.inner, clone);
172
173        Box::pin(async move {
174            let requested_tenant_id = match config.strategy {
175                TenantStrategy::Header => req
176                    .headers()
177                    .get(&config.header_name)
178                    .and_then(|v| v.to_str().ok())
179                    .map(|s| s.to_string()),
180                TenantStrategy::Subdomain => req
181                    .headers()
182                    .get(axum::http::header::HOST)
183                    .and_then(|v| v.to_str().ok())
184                    .and_then(|host| {
185                        let sub = extract_subdomain(host);
186                        if sub.is_none() {
187                            config.domain_fallback.clone()
188                        } else {
189                            sub
190                        }
191                    }),
192                TenantStrategy::Parameter => {
193                    let query = req.uri().query().unwrap_or("");
194                    serde_urlencoded::from_str::<std::collections::HashMap<String, String>>(query)
195                        .ok()
196                        .and_then(|params| params.get(&config.parameter_name).cloned())
197                }
198            };
199
200            let selected_context = req
201                .extensions()
202                .get::<crate::security::TenantMembership>()
203                .and_then(|membership| match requested_tenant_id.as_deref() {
204                    Some(requested) => membership.select(requested).ok(),
205                    None => membership.default_context(),
206                });
207
208            let Some(selected_context) = selected_context else {
209                let mut response = axum::http::Response::new(ResBody::default());
210                *response.status_mut() = axum::http::StatusCode::FORBIDDEN;
211                return Ok(response);
212            };
213            let tenant_id = selected_context.tenant_id.clone();
214            req.extensions_mut().insert(selected_context);
215            let cell = RefCell::new(Some(tenant_id));
216            TENANT_CONTEXT
217                .scope(cell, async move { inner.call(req).await })
218                .await
219        })
220    }
221}
222
223/// Fluent helper to return a Tower Middleware Layer using TenantConfig
224pub fn tenant_layer(config: TenantConfig) -> TenantLayer {
225    TenantLayer::new(config)
226}
227
228#[cfg(test)]
229#[allow(clippy::unwrap_used, clippy::expect_used)]
230mod tests {
231    use super::*;
232
233    #[test]
234    fn test_extract_subdomain() {
235        assert_eq!(
236            extract_subdomain("tenant1.example.com"),
237            Some("tenant1".to_string())
238        );
239        assert_eq!(
240            extract_subdomain("tenant-a.app.co.uk"),
241            Some("tenant-a".to_string())
242        );
243        assert_eq!(extract_subdomain("localhost:3000"), None);
244        assert_eq!(extract_subdomain("127.0.0.1"), None);
245    }
246
247    #[test]
248    fn test_tenant_config_builder() {
249        let config = TenantConfig::new(TenantStrategy::Header)
250            .with_header_name("X-Custom-Tenant")
251            .with_parameter_name("t_id")
252            .with_domain_fallback("default");
253
254        assert_eq!(config.strategy, TenantStrategy::Header);
255        assert_eq!(config.header_name, "X-Custom-Tenant");
256        assert_eq!(config.parameter_name, "t_id");
257        assert_eq!(config.domain_fallback, Some("default".to_string()));
258    }
259
260    #[tokio::test]
261    async fn test_task_local_storage() {
262        let cell = RefCell::new(Some("tenant123".to_string()));
263
264        TENANT_CONTEXT
265            .scope(cell, async {
266                assert_eq!(current_tenant_id(), Some("tenant123".to_string()));
267
268                // Set dynamic value mid-request
269                set_tenant_id(Some("super-tenant".to_string()));
270                assert_eq!(current_tenant_id(), Some("super-tenant".to_string()));
271
272                set_tenant_id(None);
273                assert_eq!(current_tenant_id(), None);
274            })
275            .await;
276
277        // Outside scope, it should return None
278        assert_eq!(current_tenant_id(), None);
279    }
280    #[tokio::test]
281    async fn test_current_tenant_id_uninitialized() {
282        assert_eq!(current_tenant_id(), None);
283    }
284
285    #[tokio::test]
286    async fn test_current_tenant_id_initialized() {
287        let cell = RefCell::new(Some("tenant-456".to_string()));
288        TENANT_CONTEXT
289            .scope(cell, async {
290                assert_eq!(current_tenant_id(), Some("tenant-456".to_string()));
291            })
292            .await;
293    }
294
295    #[tokio::test]
296    async fn test_tenant_layer_header_and_query() {
297        use axum::body::Body;
298        use axum::http::{Request, StatusCode};
299        use axum::response::IntoResponse;
300        use axum::routing::get;
301        use tower::ServiceExt;
302
303        async fn handler() -> impl IntoResponse {
304            let tenant = current_tenant_id().unwrap_or_else(|| "none".to_string());
305            (StatusCode::OK, tenant)
306        }
307
308        // 1. Header strategy
309        let config_header = TenantConfig::new(TenantStrategy::Header);
310        let app_header = axum::Router::new()
311            .route("/test", get(handler))
312            .layer(tenant_layer(config_header))
313            .layer(axum::Extension(
314                crate::security::TenantMembership::try_new(["acme-corp"]).unwrap(),
315            ));
316
317        let req = Request::builder()
318            .uri("/test")
319            .header("X-Tenant-ID", "acme-corp")
320            .body(Body::empty())
321            .unwrap();
322
323        let resp = app_header.oneshot(req).await.unwrap();
324        assert_eq!(resp.status(), StatusCode::OK);
325        let body = axum::body::to_bytes(resp.into_body(), 1000).await.unwrap();
326        assert_eq!(String::from_utf8(body.to_vec()).unwrap(), "acme-corp");
327
328        // 2. Query param strategy
329        let config_param = TenantConfig::new(TenantStrategy::Parameter);
330        let app_param = axum::Router::new()
331            .route("/test", get(handler))
332            .layer(tenant_layer(config_param))
333            .layer(axum::Extension(
334                crate::security::TenantMembership::try_new(["beta-inc"]).unwrap(),
335            ));
336
337        let req_param = Request::builder()
338            .uri("/test?tenant_id=beta-inc")
339            .body(Body::empty())
340            .unwrap();
341
342        let resp_param = app_param.oneshot(req_param).await.unwrap();
343        assert_eq!(resp_param.status(), StatusCode::OK);
344        let body_param = axum::body::to_bytes(resp_param.into_body(), 1000)
345            .await
346            .unwrap();
347        assert_eq!(String::from_utf8(body_param.to_vec()).unwrap(), "beta-inc");
348
349        // A client cannot switch to a tenant absent from authenticated claims.
350        let config_rejected = TenantConfig::new(TenantStrategy::Header);
351        let rejected_app = axum::Router::new()
352            .route("/test", get(handler))
353            .layer(tenant_layer(config_rejected))
354            .layer(axum::Extension(
355                crate::security::TenantMembership::try_new(["acme-corp"]).unwrap(),
356            ));
357        let rejected = rejected_app
358            .oneshot(
359                Request::builder()
360                    .uri("/test")
361                    .header("X-Tenant-ID", "other-tenant")
362                    .body(Body::empty())
363                    .unwrap(),
364            )
365            .await
366            .unwrap();
367        assert_eq!(rejected.status(), StatusCode::FORBIDDEN);
368    }
369
370    /// Inner service that accepts a call only on the instance that was
371    /// readied, like tower's `ConcurrencyLimit`, `RateLimit` and `Buffer`.
372    #[derive(Default)]
373    struct ReadiedOnly {
374        ready: bool,
375    }
376
377    impl Clone for ReadiedOnly {
378        fn clone(&self) -> Self {
379            Self::default()
380        }
381    }
382
383    impl tower_service::Service<axum::http::Request<axum::body::Body>> for ReadiedOnly {
384        type Response = axum::http::Response<axum::body::Body>;
385        type Error = &'static str;
386        type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
387
388        fn poll_ready(
389            &mut self,
390            _cx: &mut std::task::Context<'_>,
391        ) -> std::task::Poll<Result<(), Self::Error>> {
392            self.ready = true;
393            std::task::Poll::Ready(Ok(()))
394        }
395
396        fn call(&mut self, _req: axum::http::Request<axum::body::Body>) -> Self::Future {
397            let readied = std::mem::take(&mut self.ready);
398            std::future::ready(if readied {
399                Ok(axum::http::Response::new(axum::body::Body::empty()))
400            } else {
401                Err("called without poll_ready")
402            })
403        }
404    }
405
406    #[tokio::test]
407    async fn tenant_service_calls_the_inner_service_it_readied() {
408        use tower::ServiceExt;
409        use tower_layer::Layer;
410
411        let service =
412            tenant_layer(TenantConfig::new(TenantStrategy::Header)).layer(ReadiedOnly::default());
413        let mut request = axum::http::Request::builder()
414            .uri("/")
415            .header("X-Tenant-ID", "acme-corp")
416            .body(axum::body::Body::empty())
417            .unwrap();
418        request
419            .extensions_mut()
420            .insert(crate::security::TenantMembership::try_new(["acme-corp"]).unwrap());
421
422        let status = service
423            .oneshot(request)
424            .await
425            .map(|response| response.status());
426        assert_eq!(status, Ok(axum::http::StatusCode::OK));
427    }
428}
rullst_core/multitenant.rs

rullst_core/
multitenant.rs
