diff --git a/e2e/src/tests/tests.rs b/e2e/src/tests/tests.rs index 284caf5e8..993118d75 100644 --- a/e2e/src/tests/tests.rs +++ b/e2e/src/tests/tests.rs @@ -1409,6 +1409,131 @@ pub fn try_head() -> State { State::Success } +fn try_wildcard() -> State { + use sozu_command_lib::proto::command::{PathRule, RulePosition}; + let front_address = create_local_address(); + + let (config, listeners, state) = Worker::empty_config(); + let mut worker = Worker::start_new_worker("WLD_CRD", config, &listeners, state); + worker.send_proxy_request( + RequestType::AddHttpListener( + ListenerBuilder::new_http(front_address) + .to_http(None) + .unwrap(), + ) + .into(), + ); + worker.send_proxy_request( + RequestType::ActivateListener(ActivateListener { + address: front_address.to_string(), + proxy: ListenerType::Http.into(), + from_scm: false, + }) + .into(), + ); + + worker.send_proxy_request( + RequestType::AddCluster(Worker::default_cluster("cluster_0", false)).into(), + ); + worker.send_proxy_request( + RequestType::AddCluster(Worker::default_cluster("cluster_1", false)).into(), + ); + + worker.send_proxy_request( + RequestType::AddHttpFrontend(RequestHttpFrontend { + cluster_id: Some("cluster_0".to_string()), + address: front_address.to_string(), + hostname: String::from("*.sozu.io"), + path: PathRule::prefix(String::from("")), + position: RulePosition::Tree.into(), + ..Default::default() + }) + .into(), + ); + + let back_address: SocketAddr = create_local_address(); + worker.send_proxy_request( + RequestType::AddBackend(Worker::default_backend( + "cluster_0", + "cluster_0-0", + back_address.to_string(), + None, + )) + .into(), + ); + worker.read_to_last(); + + let mut backend0 = SyncBackend::new( + "BACKEND_0", + back_address, + http_ok_response(format!("pong0")), + ); + + let mut client = Client::new( + "client", + front_address, + http_request( + "POST", + "/api", + format!("ping"), + "www.sozu.io", + ), + ); + + backend0.connect(); + client.connect(); + client.send(); + let accepted = backend0.accept(0); + assert!(accepted); + let request = backend0.receive(0); + println!("request: {request:?}"); + backend0.send(0); + let response = client.receive(); + println!("response: {response:?}"); + + worker.send_proxy_request( + RequestType::AddHttpFrontend(RequestHttpFrontend { + cluster_id: Some("cluster_1".to_string()), + address: front_address.to_string(), + hostname: String::from("*.sozu.io"), + path: PathRule::prefix(String::from("/api")), + position: RulePosition::Tree.into(), + ..Default::default() + }) + .into(), + ); + let back_address: SocketAddr = create_local_address(); + worker.send_proxy_request( + RequestType::AddBackend(Worker::default_backend( + "cluster_1", + "cluster_1-0", + back_address.to_string(), + None, + )) + .into(), + ); + + let mut backend1 = SyncBackend::new( + "BACKEND_1", + back_address, + http_ok_response(format!("pong1")), + ); + + worker.read_to_last(); + + backend1.connect(); + + client.send(); + let accepted = backend1.accept(0); + assert!(accepted); + let request = backend1.receive(0); + println!("request: {request:?}"); + backend1.send(0); + let response = client.receive(); + println!("response: {response:?}"); + + State::Success +} #[test] fn test_sync() { @@ -1580,3 +1705,11 @@ fn test_head() { State::Success ); } + +#[test] +fn test_wildcard() { + assert_eq!( + repeat_until_error_or(1, "Hostname with wildcard", try_wildcard), + State::Success + ); +} diff --git a/lib/src/router/mod.rs b/lib/src/router/mod.rs index e0c9fc2c7..afe619a4d 100644 --- a/lib/src/router/mod.rs +++ b/lib/src/router/mod.rs @@ -59,7 +59,7 @@ impl Router { if let Some((_, path_rules)) = self.tree.lookup(hostname, true) { let mut prefix_length = 0; - let mut res = None; + let mut route = None; for (rule, method_rule, cluster_id) in path_rules { match rule.matches(path) { @@ -68,7 +68,7 @@ impl Router { MethodRuleResult::Equals => return Some(cluster_id.clone()), MethodRuleResult::All => { prefix_length = path.len(); - res = Some(cluster_id); + route = Some(cluster_id); } MethodRuleResult::None => {} } @@ -79,11 +79,11 @@ impl Router { // FIXME: the rule order will be important here MethodRuleResult::Equals => { prefix_length = size; - res = Some(cluster_id); + route = Some(cluster_id); } MethodRuleResult::All => { prefix_length = size; - res = Some(cluster_id); + route = Some(cluster_id); } MethodRuleResult::None => {} } @@ -93,7 +93,7 @@ impl Router { } } - if let Some(cluster_id) = res { + if let Some(cluster_id) = route { return Some(cluster_id.clone()); } } @@ -493,16 +493,16 @@ pub enum PathRuleResult { impl PathRule { pub fn matches(&self, path: &[u8]) -> PathRuleResult { match self { - PathRule::Prefix(s) => { - if path.starts_with(s.as_bytes()) { - PathRuleResult::Prefix(s.len()) + PathRule::Prefix(prefix) => { + if path.starts_with(prefix.as_bytes()) { + PathRuleResult::Prefix(prefix.len()) } else { PathRuleResult::None } } - PathRule::Regex(r) => { + PathRule::Regex(regex) => { let start = Instant::now(); - let is_a_match = r.is_match(path); + let is_a_match = regex.is_match(path); let now = Instant::now(); time!("regex_matching_time", (now - start).whole_milliseconds()); @@ -512,8 +512,8 @@ impl PathRule { PathRuleResult::None } } - PathRule::Equals(s) => { - if path == s.as_bytes() { + PathRule::Equals(pattern) => { + if path == pattern.as_bytes() { PathRuleResult::Equals } else { PathRuleResult::None @@ -685,6 +685,126 @@ mod tests { ); } + /// [io] + /// \ + /// [sozu] + /// \ + /// [*] <- this wildcard has multiple children + /// / \ + /// (base) (api) + #[test] + fn multiple_children_on_a_wildcard() { + let mut router = Router::new(); + + assert!(router.add_tree_rule( + b"*.sozu.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("base".to_string()) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/api".as_bytes(), &Method::Get), + Some(Route::ClusterId("base".to_string())) + ); + assert!(router.add_tree_rule( + b"*.sozu.io", + &PathRule::Prefix("/api".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("api".to_string()) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/ap".as_bytes(), &Method::Get), + Some(Route::ClusterId("base".to_string())) + ); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/api".as_bytes(), &Method::Get), + Some(Route::ClusterId("api".to_string())) + ); + } + + /// [io] + /// \ + /// [sozu] <- this node has multiple children including a wildcard + /// / \ + /// (api) [*] <- this wildcard has multiple children + /// \ + /// (base) + #[test] + fn multiple_children_including_one_with_wildcard() { + let mut router = Router::new(); + + assert!(router.add_tree_rule( + b"*.sozu.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("base".to_string()) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/api".as_bytes(), &Method::Get), + Some(Route::ClusterId("base".to_string())) + ); + assert!(router.add_tree_rule( + b"api.sozu.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("api".to_string()) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/api".as_bytes(), &Method::Get), + Some(Route::ClusterId("base".to_string())) + ); + assert_eq!( + router.lookup("api.sozu.io".as_bytes(), "/api".as_bytes(), &Method::Get), + Some(Route::ClusterId("api".to_string())) + ); + } + + #[test] + fn router_insert_remove_through_regex() { + let mut router = Router::new(); + + assert!(router.add_tree_rule( + b"www./.*/.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("base".to_string()) + )); + println!("{:#?}", router.tree); + assert!(router.add_tree_rule( + b"www.doc./.*/.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())), + &Route::ClusterId("doc".to_string()) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/".as_bytes(), &Method::Get), + Some(Route::ClusterId("base".to_string())) + ); + assert_eq!( + router.lookup("www.doc.sozu.io".as_bytes(), "/".as_bytes(), &Method::Get), + Some(Route::ClusterId("doc".to_string())) + ); + assert!(router.remove_tree_rule( + b"www./.*/.io", + &PathRule::Prefix("".to_string()), + &MethodRule::new(Some("GET".to_string())) + )); + println!("{:#?}", router.tree); + assert_eq!( + router.lookup("www.sozu.io".as_bytes(), "/".as_bytes(), &Method::Get), + None + ); + assert_eq!( + router.lookup("www.doc.sozu.io".as_bytes(), "/".as_bytes(), &Method::Get), + Some(Route::ClusterId("doc".to_string())) + ); + } + #[test] fn match_router() { let mut router = Router::new(); diff --git a/lib/src/router/pattern_trie.rs b/lib/src/router/pattern_trie.rs index 40c00ed01..f63511586 100644 --- a/lib/src/router/pattern_trie.rs +++ b/lib/src/router/pattern_trie.rs @@ -28,6 +28,12 @@ fn find_last_slash(input: &[u8]) -> Option { (0..input.len()).rev().find(|&i| input[i] == b'/') } +/// Implementation of a trie tree structure. +/// In Sozu this is used to store and lookup domains recursively. +/// Each node represents a "level domain". +/// A leaf node (leftmost label) can be a wildcard, a regex pattern or a plain string. +/// Leaves also store a value associated with the complete domain. +/// For Sozu it is a list of (PathRule, MethodRule, ClusterId). See the Router strucure. #[derive(Debug)] pub struct TrieNode { key_value: Option>, @@ -81,7 +87,10 @@ impl TrieNode { } pub fn is_empty(&self) -> bool { - self.key_value.is_none() && self.wildcard.is_none() && self.children.is_empty() + self.key_value.is_none() + && self.wildcard.is_none() + && self.regexps.is_empty() + && self.children.is_empty() } pub fn insert(&mut self, key: Key, value: V) -> InsertResult { @@ -93,9 +102,9 @@ impl TrieNode { return InsertResult::Failed; } - let res = self.insert_recursive(&key, &key, value); - assert_ne!(res, InsertResult::Failed); - res + let insert_result = self.insert_recursive(&key, &key, value); + assert_ne!(insert_result, InsertResult::Failed); + insert_result } pub fn insert_recursive(&mut self, partial_key: &[u8], key: &Key, value: V) -> InsertResult { @@ -111,9 +120,9 @@ impl TrieNode { } if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) { - for t in self.regexps.iter() { + for t in self.regexps.iter_mut() { if t.0.as_str() == s { - return InsertResult::Existing; + return t.1.insert_recursive(&partial_key[..pos - 1], key, value); } } @@ -205,12 +214,28 @@ impl TrieNode { let pos = find_last_slash(&partial_key[..partial_key.len() - 1]); if let Some(pos) = pos { + if pos > 0 && partial_key[pos - 1] != b'.' { + return RemoveResult::NotFound; + } + if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) { - let len = self.regexps.len(); - // FIXME: we might have multiple entries with the same regex - self.regexps.retain(|(r, _)| r.as_str() != s); - if len > self.regexps.len() { - return RemoveResult::Ok; + if pos > 0 { + let mut remove_result = RemoveResult::NotFound; + for t in self.regexps.iter_mut() { + if t.0.as_str() == s { + if t.1.remove_recursive(&partial_key[..pos - 1]) == RemoveResult::Ok + { + remove_result = RemoveResult::Ok; + } + } + } + return remove_result; + } else { + let len = self.regexps.len(); + self.regexps.retain(|(r, _)| r.as_str() != s); + if len > self.regexps.len() { + return RemoveResult::Ok; + } } } } @@ -229,18 +254,14 @@ impl TrieNode { Some(child) => match child.remove_recursive(prefix) { RemoveResult::NotFound => return RemoveResult::NotFound, RemoveResult::Ok => { - if !child.is_empty() { - return RemoveResult::Ok; + if child.is_empty() { + self.children.remove(suffix); } + return RemoveResult::Ok; } }, None => return RemoveResult::NotFound, } - - // if we reach here, that means we called remove_recursive on a child - // and the child is now empty - self.children.remove(suffix); - RemoveResult::Ok } pub fn lookup(&self, partial_key: &[u8], accept_wildcard: bool) -> Option<&KeyValue> { @@ -261,7 +282,6 @@ impl TrieNode { Some(child) => child.lookup(prefix, accept_wildcard), None => { //println!("no child found, testing wildcard and regexps"); - //self.print(); if prefix.is_empty() && self.wildcard.is_some() && accept_wildcard { //println!("no dot, wildcard applies"); @@ -269,15 +289,15 @@ impl TrieNode { } else { //println!("there's still a subdomain, wildcard does not apply"); - for (ref r, ref child) in self.regexps.iter() { - let s = if suffix[0] == b'.' { + for (ref regexp, ref child) in self.regexps.iter() { + let suffix = if suffix[0] == b'.' { &suffix[1..] } else { suffix }; //println!("testing regexp: {} on suffix {}", r.as_str(), str::from_utf8(s).unwrap()); - if r.is_match(s) { + if regexp.is_match(suffix) { //println!("matched"); return child.lookup(prefix, accept_wildcard); } @@ -300,6 +320,30 @@ impl TrieNode { return self.key_value.as_mut(); } + if partial_key == &b"*"[..] { + return self.wildcard.as_mut(); + } + + if partial_key[partial_key.len() - 1] == b'/' { + let pos = find_last_slash(&partial_key[..partial_key.len() - 1]); + + if let Some(pos) = pos { + if pos > 0 && partial_key[pos - 1] != b'.' { + return None; + } + + if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) { + for t in self.regexps.iter_mut() { + if t.0.as_str() == s { + return t.1.lookup_mut(&partial_key[..pos - 1], accept_wildcard); + } + } + } + } + + return None; + } + let pos = find_last_dot(partial_key); let (prefix, suffix) = match pos { None => (&b""[..], partial_key), @@ -318,15 +362,15 @@ impl TrieNode { } else { //println!("there's still a subdomain, wildcard does not apply"); - for (ref r, ref mut child) in self.regexps.iter_mut() { - let s = if suffix[0] == b'.' { + for (ref regexp, ref mut child) in self.regexps.iter_mut() { + let suffix = if suffix[0] == b'.' { &suffix[1..] } else { suffix }; //println!("testing regexp: {} on suffix {}", r.as_str(), str::from_utf8(s).unwrap()); - if r.is_match(s) { + if regexp.is_match(suffix) { //println!("matched"); return child.lookup_mut(prefix, accept_wildcard); } @@ -494,6 +538,45 @@ mod tests { assert_eq!(root, root3); } + #[test] + fn insert_remove_through_regex() { + let mut root: TrieNode = TrieNode::root(); + println!("creating root:"); + root.print(); + + println!("adding (www./.*/.com, 1)"); + assert_eq!( + root.insert(Vec::from(&b"www./.*/.com"[..]), 1), + InsertResult::Ok + ); + root.print(); + println!("adding (www.doc./.*/.com, 2)"); + assert_eq!( + root.insert(Vec::from(&b"www.doc./.*/.com"[..]), 2), + InsertResult::Ok + ); + root.print(); + assert_eq!( + root.domain_lookup(&b"www.sozu.com".to_vec(), false), + Some(&(b"www./.*/.com".to_vec(), 1)) + ); + assert_eq!( + root.domain_lookup(&b"www.doc.sozu.com".to_vec(), false), + Some(&(b"www.doc./.*/.com".to_vec(), 2)) + ); + + assert_eq!( + root.domain_remove(&b"www./.*/.com".to_vec()), + RemoveResult::Ok + ); + root.print(); + assert_eq!(root.domain_lookup(&b"www.sozu.com".to_vec(), false), None); + assert_eq!( + root.domain_lookup(&b"www.doc.sozu.com".to_vec(), false), + Some(&(b"www.doc./.*/.com".to_vec(), 2)) + ); + } + #[test] fn add_child_to_leaf() { let mut root1: TrieNode = TrieNode::root();