use std::collections::{BTreeSet, HashMap};
use serde::{Deserialize, Serialize};
use crate::dns::DnsResolver;
use crate::takeover::{resolve_host, truncate_to_cap};
#[derive(Debug, Clone, PartialEq, Eq)]
struct WildcardAnswer {
cname: Option<String>,
}
const MAX_CLASSIFY_NAMES: usize = 2000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SubdomainStatus {
Live,
Dead,
Wildcard,
Unknown,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifiedSubdomain {
pub name: String,
pub status: SubdomainStatus,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub addresses: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cname: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub takeover_risk: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubdomainClassification {
pub domain: String,
pub wildcard_detected: bool,
pub subdomains: Vec<ClassifiedSubdomain>,
#[serde(default)]
pub names_skipped: usize,
}
fn cname_key(cname: Option<&str>) -> Option<String> {
cname.map(|c| c.trim_end_matches('.').to_ascii_lowercase())
}
fn classify_one(
name: String,
addresses: Vec<String>,
cname: Option<String>,
provider: Option<&'static str>,
lookup_failed: bool,
wildcard: Option<&WildcardAnswer>,
) -> ClassifiedSubdomain {
let status = if addresses.is_empty() && lookup_failed {
SubdomainStatus::Unknown
} else if addresses.is_empty() {
SubdomainStatus::Dead
} else {
match wildcard {
Some(w) if cname_key(cname.as_deref()) == cname_key(w.cname.as_deref()) => {
SubdomainStatus::Wildcard
}
_ => SubdomainStatus::Live,
}
};
let takeover_risk = match status {
SubdomainStatus::Dead => provider.map(str::to_string),
_ => None,
};
ClassifiedSubdomain {
name,
status,
addresses,
cname,
takeover_risk,
}
}
fn parent_of(name: &str) -> Option<&str> {
name.split_once('.').map(|(_, parent)| parent)
}
async fn probe_parents(
resolver: &DnsResolver,
parents: BTreeSet<String>,
concurrency: usize,
) -> HashMap<String, Option<WildcardAnswer>> {
use futures::stream::{self, StreamExt};
stream::iter(parents)
.map(|parent| async move {
let probe = format!("{}.{parent}", crate::dns::random_probe_label());
let r = resolve_host(resolver, &probe).await;
let answer = (!r.addresses.is_empty()).then_some(WildcardAnswer { cname: r.cname });
(parent, answer)
})
.buffer_unordered(concurrency)
.collect()
.await
}
pub async fn classify_subdomains(
resolver: &DnsResolver,
domain: &str,
mut names: Vec<String>,
concurrency: usize,
) -> SubdomainClassification {
use futures::stream::{self, StreamExt};
let names_skipped = truncate_to_cap(&mut names, MAX_CLASSIFY_NAMES);
let concurrency = concurrency.max(1);
let parents: BTreeSet<String> = names
.iter()
.filter_map(|n| parent_of(n))
.map(str::to_string)
.collect();
let wildcards = probe_parents(resolver, parents, concurrency).await;
let wildcard_detected = wildcards.values().any(Option::is_some);
let wildcards = &wildcards;
let mut subdomains: Vec<ClassifiedSubdomain> = stream::iter(names)
.map(|name| async move {
let r = resolve_host(resolver, &name).await;
let failed = r.lookup_error.is_some();
let wildcard = parent_of(&name)
.and_then(|p| wildcards.get(p))
.and_then(Option::as_ref);
let provider = r.provider.map(|(p, _)| p.provider);
classify_one(name, r.addresses, r.cname, provider, failed, wildcard)
})
.buffer_unordered(concurrency)
.collect()
.await;
subdomains.sort_by(|a, b| {
let rank = |s: &ClassifiedSubdomain| {
if s.takeover_risk.is_some() {
0
} else if s.status == SubdomainStatus::Live {
1
} else {
2
}
};
rank(a).cmp(&rank(b)).then_with(|| a.name.cmp(&b.name))
});
SubdomainClassification {
domain: domain.to_string(),
wildcard_detected,
subdomains,
names_skipped,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::takeover::match_provider;
fn classify(
name: String,
addresses: Vec<String>,
cname: Option<String>,
lookup_failed: bool,
wildcard: Option<&WildcardAnswer>,
) -> ClassifiedSubdomain {
let provider = cname
.as_deref()
.and_then(match_provider)
.map(|p| p.provider);
classify_one(name, addresses, cname, provider, lookup_failed, wildcard)
}
#[test]
fn classify_flags_a_provider_deeper_in_the_chain() {
let c = classify_one(
"shop.example.com".to_string(),
vec![],
Some("edge.example.net".to_string()),
Some("Heroku"),
false,
None,
);
assert_eq!(c.takeover_risk.as_deref(), Some("Heroku"));
}
#[test]
fn classify_flags_every_takeover_provider() {
let dangling = |cname: &str| {
classify(
"gone.example.com".to_string(),
vec![],
Some(cname.to_string()),
false,
None,
)
.takeover_risk
};
for p in crate::takeover::PROVIDERS {
let suffixes = p.cname_suffixes.iter().map(|s| format!("gone{s}"));
let infixes = p.cname_infixes.iter().map(|i| format!("gone{i}.example"));
for cname in suffixes.chain(infixes) {
assert_eq!(dangling(&cname).as_deref(), Some(p.provider), "{cname}");
}
}
assert_eq!(dangling("cdn.example.com"), None);
assert_eq!(dangling("github.io"), None);
}
#[test]
fn classify_dead_cname_to_provider_is_takeover_risk() {
let c = classify(
"gone.example.com".to_string(),
vec![], Some("gone.herokuapp.com".to_string()),
false,
None,
);
assert_eq!(c.status, SubdomainStatus::Dead);
assert_eq!(c.takeover_risk.as_deref(), Some("Heroku"));
}
#[test]
fn classify_failed_lookup_is_unknown_not_dangling() {
let c = classify(
"slow.example.com".to_string(),
vec![],
Some("slow.herokuapp.com".to_string()),
true,
None,
);
assert_eq!(c.status, SubdomainStatus::Unknown);
assert!(c.takeover_risk.is_none());
let c = classify(
"v6.example.com".to_string(),
vec!["2001:db8::1".to_string()],
None,
true,
None,
);
assert_eq!(c.status, SubdomainStatus::Live);
}
#[test]
fn classify_live_host_is_not_takeover_even_with_provider_cname() {
let c = classify(
"live.example.com".to_string(),
vec!["203.0.113.5".to_string()],
Some("live.herokuapp.com".to_string()),
false,
None,
);
assert_eq!(c.status, SubdomainStatus::Live);
assert!(c.takeover_risk.is_none());
}
#[test]
fn classify_under_a_wildcard_compares_presence_not_addresses() {
let wildcard = WildcardAnswer { cname: None };
let c = classify(
"anything.example.com".to_string(),
vec!["203.0.113.7".to_string()],
None,
false,
Some(&wildcard),
);
assert_eq!(c.status, SubdomainStatus::Wildcard);
let wildcard = WildcardAnswer {
cname: Some("edge.cdn.test.".to_string()),
};
let c = classify(
"x.example.com".to_string(),
vec!["203.0.113.7".to_string()],
Some("EDGE.cdn.test".to_string()),
false,
Some(&wildcard),
);
assert_eq!(c.status, SubdomainStatus::Wildcard);
}
#[test]
fn classify_a_distinct_alias_under_a_wildcard_is_live() {
let wildcard = WildcardAnswer { cname: None };
let c = classify(
"shop.example.com".to_string(),
vec!["203.0.113.7".to_string()],
Some("shops.myshopify.com".to_string()),
false,
Some(&wildcard),
);
assert_eq!(c.status, SubdomainStatus::Live);
let c = classify(
"real.example.com".to_string(),
vec!["203.0.113.7".to_string()],
None,
false,
None,
);
assert_eq!(c.status, SubdomainStatus::Live);
}
#[tokio::test]
async fn wildcards_below_the_apex_are_detected() {
use crate::dns::test_support::{
a_rdata, mock_dns_resolver_default, spawn_mock_dns_fn, MockReply,
};
use hickory_resolver::proto::rr::RecordType as WireType;
let mut rotation = 10u8;
let port = spawn_mock_dns_fn(move |qname, qtype| {
let in_dev = qname.ends_with(".dev.seer.test");
match qtype {
WireType::A if in_dev => {
rotation += 1;
MockReply::Answer(vec![a_rdata([192, 0, 2, rotation])])
}
WireType::A if qname == "www.seer.test" => {
MockReply::Answer(vec![a_rdata([192, 0, 2, 1])])
}
_ if in_dev || qname == "www.seer.test" => MockReply::NoData,
_ => MockReply::NxDomain,
}
})
.await;
let report = classify_subdomains(
&mock_dns_resolver_default(port),
"seer.test",
vec!["api.dev.seer.test".to_string(), "www.seer.test".to_string()],
4,
)
.await;
assert!(report.wildcard_detected);
let status = |name: &str| {
report
.subdomains
.iter()
.find(|s| s.name == name)
.map(|s| s.status)
};
assert_eq!(status("api.dev.seer.test"), Some(SubdomainStatus::Wildcard));
assert_eq!(status("www.seer.test"), Some(SubdomainStatus::Live));
}
#[test]
fn classify_cap_keeps_first_n_and_reports_skipped() {
let over = MAX_CLASSIFY_NAMES + 37;
let mut names: Vec<String> = (0..over).map(|i| format!("h{i}.example.com")).collect();
let skipped = truncate_to_cap(&mut names, MAX_CLASSIFY_NAMES);
assert_eq!(names.len(), MAX_CLASSIFY_NAMES);
assert_eq!(skipped, 37);
assert_eq!(names[0], "h0.example.com");
}
#[test]
fn classify_cap_is_noop_under_limit() {
let mut names: Vec<String> = (0..10).map(|i| format!("h{i}.example.com")).collect();
let skipped = truncate_to_cap(&mut names, MAX_CLASSIFY_NAMES);
assert_eq!(names.len(), 10);
assert_eq!(skipped, 0);
}
}