use super::traits::{Identifier, IdentifierExtractor, RequestContext};
use std::net::IpAddr;
pub struct CompositeExtractor {
extractors: Vec<Box<dyn IdentifierExtractor>>,
fallback_to_default: bool,
}
#[derive(Default)]
pub struct CompositeExtractorBuilder {
extractors: Vec<Box<dyn IdentifierExtractor>>,
fallback_to_default: bool,
}
impl CompositeExtractorBuilder {
pub fn new() -> Self {
Self {
extractors: Vec::new(),
fallback_to_default: false,
}
}
pub fn add_extractor(mut self, extractor: Box<dyn IdentifierExtractor>) -> Self {
self.extractors.push(extractor);
self
}
pub fn with_fallback(mut self, fallback: bool) -> Self {
self.fallback_to_default = fallback;
self
}
pub fn build(self) -> CompositeExtractor {
CompositeExtractor::with_dependencies(self.extractors, self.fallback_to_default)
}
}
impl CompositeExtractor {
pub fn new(extractors: Vec<Box<dyn IdentifierExtractor>>, fallback_to_default: bool) -> Self {
Self {
extractors,
fallback_to_default,
}
}
pub fn builder() -> CompositeExtractorBuilder {
CompositeExtractorBuilder::new()
}
pub fn with_dependencies(
extractors: Vec<Box<dyn IdentifierExtractor>>,
fallback_to_default: bool,
) -> Self {
Self {
extractors,
fallback_to_default,
}
}
pub fn add_extractor(mut self, extractor: Box<dyn IdentifierExtractor>) -> Self {
self.extractors.push(extractor);
self
}
pub fn with_fallback(mut self, fallback: bool) -> Self {
self.fallback_to_default = fallback;
self
}
}
impl IdentifierExtractor for CompositeExtractor {
fn extract(&self, context: &RequestContext) -> Option<Identifier> {
for extractor in &self.extractors {
if let Some(identifier) = extractor.extract(context) {
return Some(identifier);
}
}
if self.fallback_to_default {
if let Some(client_ip) = &context.client_ip {
if client_ip.parse::<IpAddr>().is_ok() {
return Some(Identifier::Ip(client_ip.clone()));
}
}
}
None
}
fn name(&self) -> &str {
"CompositeExtractor"
}
}
#[cfg(test)]
#[allow(clippy::items_after_test_module)]
mod tests {
use super::*;
use crate::matchers::{IpExtractor, UserIdExtractor};
#[test]
fn test_composite_builder() {
let extractor = CompositeExtractor::builder()
.add_extractor(Box::new(UserIdExtractor::from_header("X-User-Id")))
.add_extractor(Box::new(IpExtractor::from_header("X-Forwarded-For")))
.with_fallback(true)
.build();
assert_eq!(extractor.name(), "CompositeExtractor");
}
#[test]
fn test_composite_fallback_to_client_ip() {
let extractor = CompositeExtractor::new(vec![], true);
let ctx = RequestContext::new().with_client_ip("10.0.0.1");
assert_eq!(
extractor.extract(&ctx),
Some(Identifier::Ip("10.0.0.1".into()))
);
}
#[test]
fn test_composite_fallback_no_ip() {
let extractor = CompositeExtractor::new(vec![], true);
let ctx = RequestContext::new();
assert_eq!(extractor.extract(&ctx), None);
}
#[test]
fn test_composite_fallback_invalid_ip() {
let extractor = CompositeExtractor::new(vec![], true);
let ctx = RequestContext::new().with_client_ip("not-an-ip");
assert_eq!(extractor.extract(&ctx), None);
}
#[test]
fn test_composite_with_dependencies() {
let extractor = CompositeExtractor::with_dependencies(
vec![Box::new(UserIdExtractor::from_header("X-User-Id"))],
false,
);
let ctx = RequestContext::new().with_header("X-User-Id", "u1");
assert_eq!(
extractor.extract(&ctx),
Some(Identifier::UserId("u1".into()))
);
}
#[test]
fn test_composite_add_extractor() {
let extractor = CompositeExtractor::new(vec![], false)
.add_extractor(Box::new(UserIdExtractor::from_header("X-User-Id")));
let ctx = RequestContext::new().with_header("X-User-Id", "added");
assert_eq!(
extractor.extract(&ctx),
Some(Identifier::UserId("added".into()))
);
}
#[test]
fn test_composite_with_fallback() {
let extractor = CompositeExtractor::new(vec![], false).with_fallback(true);
assert_eq!(extractor.name(), "CompositeExtractor");
}
#[test]
fn test_composite_fallback_after_extractor_failure() {
let extractor = CompositeExtractor::new(
vec![Box::new(UserIdExtractor::from_header("X-User-Id"))],
true,
);
let ctx = RequestContext::new()
.with_header("X-Other-Header", "value")
.with_client_ip("10.0.0.2");
assert_eq!(
extractor.extract(&ctx),
Some(Identifier::Ip("10.0.0.2".into()))
);
}
#[test]
fn test_composite_no_fallback_after_extractor_failure() {
let extractor = CompositeExtractor::new(
vec![Box::new(UserIdExtractor::from_header("X-User-Id"))],
false,
);
let ctx = RequestContext::new()
.with_header("X-Other-Header", "value")
.with_client_ip("10.0.0.3");
assert_eq!(extractor.extract(&ctx), None);
}
#[test]
fn test_composite_first_extractor_wins() {
let extractor = CompositeExtractor::new(
vec![
Box::new(UserIdExtractor::from_header("X-User-Id")),
Box::new(IpExtractor::from_header("X-Forwarded-For")),
],
false,
);
let ctx = RequestContext::new()
.with_header("X-User-Id", "winner")
.with_header("X-Forwarded-For", "1.2.3.4");
assert_eq!(
extractor.extract(&ctx),
Some(Identifier::UserId("winner".into()))
);
}
#[test]
fn test_composite_builder_default_no_fallback() {
let extractor = CompositeExtractor::builder()
.add_extractor(Box::new(UserIdExtractor::from_header("X-User-Id")))
.build();
let ctx = RequestContext::new().with_client_ip("10.0.0.4");
assert_eq!(extractor.extract(&ctx), None);
}
}