-
-
Notifications
You must be signed in to change notification settings - Fork 120
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
124 additions
and
258 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,134 +0,0 @@ | ||
| use std::collections::HashMap; | ||
| use std::net::IpAddr; | ||
| use wildmatch::WildMatch; | ||
| use crate::dns::DNSEntry; | ||
|
|
||
| /// Local override table that can be used instead of using /etc/hosts or similar 3rd party dns system. | ||
| pub struct LocalOverrideTable { | ||
| /// Entries in the local override table. First iteration of the resolver: a simple list that | ||
| /// will be queried O(n). Later iterations will use a more efficient data structure like a tree | ||
| /// per domain part (e.g. com, org, net, etc.) | ||
| entries: Vec<DNSEntry>, | ||
| round_robin_ptrs: HashMap<String, usize>, | ||
| } | ||
|
|
||
| impl Default for LocalOverrideTable { | ||
| fn default() -> LocalOverrideTable { | ||
| LocalOverrideTable { | ||
| entries: Vec::new(), | ||
| round_robin_ptrs: HashMap::new(), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl LocalOverrideTable { | ||
|
|
||
| pub fn resolve(&mut self, domain: &str) -> Option<IpAddr> { | ||
| let mut ips = Vec::new(); | ||
| let mut wildcard_ips = Vec::new(); | ||
|
|
||
| // For now it's ok to iterate the entries list, but later we should use a more efficient | ||
| // data structure. | ||
| for entry in &self.entries { | ||
| // We separate wildcard matches and specific matches. | ||
| if WildMatch::new(&entry.domain_name).matches(domain) { | ||
| wildcard_ips.push(entry.ip); | ||
| } | ||
| if entry.domain_name == domain { | ||
| ips.push(entry.ip); | ||
| } | ||
| } | ||
|
|
||
| // if we haven't found a specific domain match, we use the wildcard domain, otherwise | ||
| // the specific matches take precedence. | ||
| if ips.len() == 0 { | ||
| ips = wildcard_ips | ||
| } | ||
|
|
||
| // If there is only one ip for a domain, return it | ||
| if ips.len() == 1 { | ||
| return Some(*ips.get(0).unwrap()); | ||
| } | ||
|
|
||
| // If there are multiple ips for a domain, use round robin to fetch the next one | ||
| if ips.len() > 1 { | ||
| let mut rr_ptr = *self.round_robin_ptrs.get(domain).unwrap_or(&0); | ||
| if rr_ptr >= ips.len() { | ||
| rr_ptr = 0; | ||
| } | ||
|
|
||
| let ip = ips.get(rr_ptr).expect("Invalid round robin pointer"); | ||
| self.round_robin_ptrs.insert(domain.to_string(), rr_ptr + 1 % ips.len()); | ||
|
|
||
| return Some(*ip); | ||
| } | ||
|
|
||
| None | ||
| } | ||
| } | ||
|
|
||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use core::str::FromStr; | ||
| use crate::dns::DnsResolver; | ||
| use super::*; | ||
|
|
||
| #[test] | ||
| fn test_local_override() { | ||
| let mut resolver = DnsResolver::new(); | ||
| resolver.local_override_enabled = true; | ||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "example.com".to_string(), | ||
| ip: IpAddr::from_str("1.2.3.4").unwrap(), | ||
| }); | ||
|
|
||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "foo.example.com".to_string(), | ||
| ip: IpAddr::from_str("2.3.4.5").unwrap(), | ||
| }); | ||
|
|
||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "*.wilcard.com".to_string(), | ||
| ip: IpAddr::from_str("6.6.6.6").unwrap(), | ||
| }); | ||
|
|
||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "specific.wilcard.com".to_string(), | ||
| ip: IpAddr::from_str("8.8.8.8").unwrap(), | ||
| }); | ||
|
|
||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "ipv6.com".to_string(), | ||
| ip: IpAddr::from_str("2002::1").unwrap(), | ||
| }); | ||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "ipv6.com".to_string(), | ||
| ip: IpAddr::from_str("2002::2").unwrap(), | ||
| }); | ||
| resolver.local_override.entries.push(DNSEntry { | ||
| domain_name: "ipv6.com".to_string(), | ||
| ip: IpAddr::from_str("200.200.200.200").unwrap(), | ||
| }); | ||
|
|
||
| // Simple resolve | ||
| assert_eq!(IpAddr::from_str("1.2.3.4").unwrap(), resolver.resolve("example.com").unwrap()); | ||
| assert!(resolver.resolve("xample.com").is_err()); | ||
| assert!(resolver.resolve("com").is_err()); | ||
| assert!(resolver.resolve("example").is_err()); | ||
|
|
||
| // Wildcard | ||
| assert_eq!(IpAddr::from_str("8.8.8.8").unwrap(), resolver.resolve("specific.wilcard.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("6.6.6.6").unwrap(), resolver.resolve("something.wilcard.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("6.6.6.6").unwrap(), resolver.resolve("foobar.wilcard.com").unwrap()); | ||
| assert!(resolver.resolve("wilcard.com").is_err()); | ||
|
|
||
| // round robin | ||
| assert_eq!(IpAddr::from_str("2002::1").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("2002::2").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("200.200.200.200").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("2002::1").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("2002::2").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| assert_eq!(IpAddr::from_str("200.200.200.200").unwrap(), resolver.resolve("ipv6.com").unwrap()); | ||
| } | ||
| } | ||
Oops, something went wrong.