fix: prevent 413 errors because of large org auth cookies

crates/service_utils/src/middlewares/auth_n/authentication.rs

@@ -1,7 +1,7 @@
 use std::fmt::Display;
 
 use actix_web::{
-    HttpRequest, HttpResponse, Scope,
+    HttpRequest, HttpResponse, HttpResponseBuilder, Scope,
     cookie::{Cookie, time::Duration},
     http::header,
     web::Path,
@@ -32,6 +32,91 @@ impl Display for Login {
     }
 }
 
+const MAX_ORG_COOKIES: usize = 2;
+const ORG_COOKIE_PREFIX: &str = "org_";
+const ORGS_ORDER_COOKIE: &str = "orgs_order";
+
+fn parse_org_ids_from_cookie_header(req: &HttpRequest) -> Vec<String> {
+    req.headers()
+        .get(header::COOKIE)
+        .and_then(|v| v.to_str().ok())
+        .map(|header_val| {
+            header_val
+                .split(';')
+                .filter_map(|pair| {
+                    let pair = pair.trim();
+                    let eq = pair.find('=')?;
+                    let name = &pair[..eq];
+                    name.strip_prefix(ORG_COOKIE_PREFIX)
+                        .map(|id| id.to_string())
+                })
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
+pub(super) fn add_org_cookie_with_eviction(
+    req: &HttpRequest,
+    login_type: &Login,
+    token: String,
+    cookie_path: String,
+    mut builder: HttpResponseBuilder,
+) -> HttpResponseBuilder {
+    let Login::Org(new_org_id) = login_type else {
+        return builder;
+    };
+
+    let mut order: Vec<String> = req
+        .cookie(ORGS_ORDER_COOKIE)
+        .and_then(|c| {
+            let v: Vec<String> = c
+                .value()
+                .split(',')
+                .filter(|s| !s.is_empty())
+                .map(|s| s.to_string())
+                .collect();
+            (!v.is_empty()).then_some(v)
+        })
+        .unwrap_or_else(|| {
+            log::warn!(
+                "orgs_order cookie missing, scanning request cookies for org cookies"
+            );
+            parse_org_ids_from_cookie_header(req)
+        });
+
+    order.retain(|id| *id != *new_org_id);
+    order.push(new_org_id.clone());
+
+    while order.len() > MAX_ORG_COOKIES {
+        let evicted_id = order.remove(0);
+        let evict_cookie = Cookie::build(format!("{ORG_COOKIE_PREFIX}{evicted_id}"), "")
+            .path(cookie_path.clone())
+            .http_only(true)
+            .secure(true)
+            .max_age(Duration::seconds(0))
+            .finish();
+        builder.cookie(evict_cookie);
+    }
+
+    let new_cookie = Cookie::build(login_type.to_string(), token)
+        .path(cookie_path.clone())
+        .http_only(true)
+        .secure(true)
+        .max_age(Duration::days(1))
+        .finish();
+    builder.cookie(new_cookie);
+
+    let order_cookie = Cookie::build(ORGS_ORDER_COOKIE.to_string(), order.join(","))
+        .path(cookie_path)
+        .http_only(true)
+        .secure(true)
+        .max_age(Duration::days(1))
+        .finish();
+    builder.cookie(order_cookie);
+
+    builder
+}
+
 pub trait Authenticator: Sync + Send {
     fn routes(&self) -> Scope;
 
@@ -63,7 +148,7 @@ pub trait Authenticator: Sync + Send {
 
     fn switch_organisation<'a>(
         &'a self,
-        req: &HttpRequest,
+        req: &'a HttpRequest,
         path: &Path<SwitchOrgParams>,
     ) -> LocalBoxFuture<'a, HttpResponse> {
         let login_type = Login::Org(path.organisation_id.clone());
@@ -76,21 +161,18 @@ pub trait Authenticator: Sync + Send {
 
         Box::pin(async move {
             match user_token_future.await {
-                Ok(token) => {
-                    let cookie = Cookie::build(login_type.to_string(), token)
-                        .path(cookie_path)
-                        .http_only(true)
-                        .secure(true)
-                        .max_age(Duration::days(1))
-                        .finish();
-                    HttpResponse::Found()
-                        .cookie(cookie)
-                        .insert_header((
-                            header::LOCATION,
-                            format!("{prefix}/admin/{org_id}/workspaces"),
-                        ))
-                        .finish()
-                }
+                Ok(token) => add_org_cookie_with_eviction(
+                    req,
+                    &login_type,
+                    token,
+                    cookie_path,
+                    HttpResponse::Found(),
+                )
+                .insert_header((
+                    header::LOCATION,
+                    format!("{prefix}/admin/{org_id}/workspaces"),
+                ))
+                .finish(),
                 Err(resp) => resp,
             }
         })

crates/service_utils/src/middlewares/auth_n/oidc/saas_authenticator.rs

@@ -2,7 +2,6 @@ use std::sync::Arc;
 
 use actix_web::{
     HttpRequest, HttpResponse,
-    cookie::{Cookie, time::Duration},
     error::{ErrorBadRequest, ErrorInternalServerError},
     http::header,
     web::{self, Data, Json, get, resource},
@@ -20,7 +19,7 @@ use crate::{
     extensions::HttpRequestExt,
     helpers::get_from_env_unsafe,
     middlewares::auth_n::{
-        authentication::{Authenticator, Login},
+        authentication::{Authenticator, Login, add_org_cookie_with_eviction},
         oidc::{
             OIDCAuthenticator,
             types::{
@@ -201,16 +200,16 @@ impl SaasOIDCAuthenticator {
             self.generate_org_user(&request, &org_id, &login_type)
                 .await
                 .and_then(|token| {
-                    let cookie = Cookie::build(login_type.to_string(), token)
-                        .path(self.get_cookie_path())
-                        .http_only(true)
-                        .secure(true)
-                        .max_age(Duration::days(1))
-                        .finish();
-                    Err(HttpResponse::Found()
-                        .cookie(cookie)
-                        .insert_header((header::LOCATION, request.path().to_string()))
-                        .finish())
+                    let response = add_org_cookie_with_eviction(
+                        &request,
+                        &login_type,
+                        token,
+                        self.get_cookie_path(),
+                        HttpResponse::Found(),
+                    )
+                    .insert_header((header::LOCATION, request.path().to_string()))
+                    .finish();
+                    Err(response)
                 })
         }
     }