use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use percent_encoding::percent_decode_str;
use crate::error::Error;
pub const COMMON_DOMAIN_COOKIE_NAME: &str = "_saml_idp";
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CommonDomainCookie {
pub entity_ids: Vec<String>,
}
impl CommonDomainCookie {
pub fn parse(cookie_value: &str) -> Result<Self, Error> {
let decoded = percent_decode_str(cookie_value)
.decode_utf8()
.map_err(|_utf8_err| Error::CommonDomainCookieMalformed {
reason: "cookie value is not valid UTF-8 after percent-decoding",
})?;
let mut entity_ids = Vec::new();
for token in decoded.split(' ') {
if token.is_empty() {
continue;
}
let bytes = BASE64_STANDARD.decode(token).map_err(|_b64_err| {
Error::CommonDomainCookieMalformed {
reason: "cookie entry is not valid base64",
}
})?;
let entity_id = String::from_utf8(bytes).map_err(|_utf8_err| {
Error::CommonDomainCookieMalformed {
reason: "cookie entry does not decode to UTF-8",
}
})?;
if entity_id.is_empty() {
return Err(Error::CommonDomainCookieMalformed {
reason: "cookie entry decodes to an empty entityID",
});
}
entity_ids.push(entity_id);
}
Ok(Self { entity_ids })
}
pub fn most_recent(&self) -> Option<&str> {
self.entity_ids.last().map(String::as_str)
}
pub fn record(&mut self, entity_id: &str) {
self.entity_ids.retain(|existing| existing != entity_id);
self.entity_ids.push(entity_id.to_owned());
}
pub fn to_cookie_value(&self) -> String {
let encoded: Vec<String> = self
.entity_ids
.iter()
.map(|id| BASE64_STANDARD.encode(id.as_bytes()))
.collect();
encoded.join("%20")
}
pub fn truncate_to_fit(&mut self, max_value_len: usize) {
while !self.entity_ids.is_empty() && self.to_cookie_value().len() > max_value_len {
self.entity_ids.remove(0);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_value_parses_to_empty_cookie() {
let cookie = CommonDomainCookie::parse("").unwrap();
assert!(cookie.entity_ids.is_empty());
assert_eq!(cookie.most_recent(), None);
assert_eq!(cookie.to_cookie_value(), "");
}
#[test]
fn roundtrip_preserves_order_and_most_recent() {
let mut cookie = CommonDomainCookie::default();
cookie.record("https://idp-a.example.com");
cookie.record("https://idp-b.example.com");
let value = cookie.to_cookie_value();
assert!(value.contains("%20"), "separator must be encoded: {value}");
let reparsed = CommonDomainCookie::parse(&value).unwrap();
assert_eq!(
reparsed.entity_ids,
vec!["https://idp-a.example.com", "https://idp-b.example.com"]
);
assert_eq!(reparsed.most_recent(), Some("https://idp-b.example.com"));
}
#[test]
fn parse_accepts_raw_space_separator() {
let raw = format!(
"{} {}",
BASE64_STANDARD.encode("https://idp-a.example.com"),
BASE64_STANDARD.encode("https://idp-b.example.com"),
);
let cookie = CommonDomainCookie::parse(&raw).unwrap();
assert_eq!(cookie.entity_ids.len(), 2);
}
#[test]
fn parse_accepts_fully_percent_encoded_value() {
let b64 = BASE64_STANDARD.encode("https://idp.example.com/saml");
let fully_encoded: String = b64
.chars()
.flat_map(|c| {
if c == '=' {
"%3D".chars().collect::<Vec<_>>()
} else {
vec![c]
}
})
.collect();
let cookie = CommonDomainCookie::parse(&fully_encoded).unwrap();
assert_eq!(cookie.entity_ids, vec!["https://idp.example.com/saml"]);
}
#[test]
fn record_moves_existing_entry_to_most_recent() {
let mut cookie = CommonDomainCookie::default();
cookie.record("https://idp-a.example.com");
cookie.record("https://idp-b.example.com");
cookie.record("https://idp-a.example.com");
assert_eq!(
cookie.entity_ids,
vec!["https://idp-b.example.com", "https://idp-a.example.com"]
);
assert_eq!(cookie.most_recent(), Some("https://idp-a.example.com"));
}
#[test]
fn base64_plus_character_survives_roundtrip() {
let entity_id = "https://idp.example.com/tenant?~~"; let b64 = BASE64_STANDARD.encode(entity_id);
assert!(b64.contains('+'), "fixture must exercise '+': {b64}");
let mut cookie = CommonDomainCookie::default();
cookie.record(entity_id);
let reparsed = CommonDomainCookie::parse(&cookie.to_cookie_value()).unwrap();
assert_eq!(reparsed.entity_ids, vec![entity_id]);
}
#[test]
fn rejects_non_base64_entry() {
let err = CommonDomainCookie::parse("not-base64!!!").unwrap_err();
assert!(matches!(err, Error::CommonDomainCookieMalformed { .. }));
}
#[test]
fn rejects_entry_decoding_to_invalid_utf8() {
let bad = BASE64_STANDARD.encode([0xFF, 0xFE, 0x00]);
let err = CommonDomainCookie::parse(&bad).unwrap_err();
assert!(matches!(
err,
Error::CommonDomainCookieMalformed {
reason: "cookie entry does not decode to UTF-8"
}
));
}
#[test]
fn truncate_to_fit_drops_oldest_first() {
let mut cookie = CommonDomainCookie::default();
for i in 0..50 {
cookie.record(&format!("https://idp-{i}.example-federation.org/saml"));
}
let full_len = cookie.to_cookie_value().len();
assert!(full_len > 1024);
cookie.truncate_to_fit(1024);
assert!(cookie.to_cookie_value().len() <= 1024);
assert_eq!(
cookie.most_recent(),
Some("https://idp-49.example-federation.org/saml")
);
assert_ne!(
cookie.entity_ids.first().map(String::as_str),
Some("https://idp-0.example-federation.org/saml")
);
}
#[test]
fn truncate_to_zero_empties_the_cookie() {
let mut cookie = CommonDomainCookie::default();
cookie.record("https://idp.example.com");
cookie.truncate_to_fit(0);
assert!(cookie.entity_ids.is_empty());
}
}