use std::net::IpAddr;
use std::sync::{Arc, Mutex};
use flowscope::Timestamp;
use flowscope::dns::{DnsMessage, DnsRcode, NameClaim, NameMap, NameMapConfig};
use crate::protocol::FlowKey;
pub struct DnsView<'a> {
pub message: &'a DnsMessage,
pub client: IpAddr,
pub server: IpAddr,
pub key: FlowKey,
pub ts: Timestamp,
}
impl DnsView<'_> {
pub fn qname(&self) -> Option<&str> {
match self.message {
DnsMessage::Query(q) => q.questions.first().map(|x| x.name.as_str()),
DnsMessage::Response(r) => r.questions.first().map(|x| x.name.as_str()),
_ => None,
}
}
pub fn rcode(&self) -> Option<DnsRcode> {
match self.message {
DnsMessage::Response(r) => Some(r.rcode),
_ => None,
}
}
pub fn is_answerless(&self) -> bool {
match self.message {
DnsMessage::Response(r) => r.answers.is_empty(),
_ => false,
}
}
pub fn is_query(&self) -> bool {
matches!(self.message, DnsMessage::Query(_))
}
pub fn community_id(&self) -> Option<String> {
flowscope::KeyFields::community_id(&self.key)
}
}
pub(crate) fn client_server(key: &FlowKey) -> (IpAddr, IpAddr) {
if key.a.port() == 53 {
(key.b.ip(), key.a.ip())
} else {
(key.a.ip(), key.b.ip())
}
}
pub type NameHandler = Arc<dyn Fn(IpAddr, &NameClaim, &mut crate::ctx::Ctx<'_>) + Send + Sync>;
pub(crate) type NameHandlers = Arc<Mutex<Vec<NameHandler>>>;
pub(crate) struct NameMapState {
pub(crate) map: NameMap,
pub(crate) scratch: Vec<(IpAddr, NameClaim)>,
}
impl NameMapState {
pub(crate) fn new(config: NameMapConfig) -> Self {
Self {
map: NameMap::with_config(config),
scratch: Vec::new(),
}
}
}
impl Default for NameMapState {
fn default() -> Self {
Self {
map: NameMap::new(),
scratch: Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use flowscope::L4Proto;
use flowscope::extract::FiveTupleKey;
#[test]
fn client_server_picks_the_port_53_side() {
let k = FiveTupleKey::new(
L4Proto::Udp,
"9.9.9.9:53".parse().unwrap(),
"10.0.0.5:40000".parse().unwrap(),
);
let (c, s) = client_server(&k);
assert_eq!(c.to_string(), "10.0.0.5");
assert_eq!(s.to_string(), "9.9.9.9");
let k = FiveTupleKey::new(
L4Proto::Udp,
"10.0.0.5:40000".parse().unwrap(),
"9.9.9.9:53".parse().unwrap(),
);
let (c, s) = client_server(&k);
assert_eq!(c.to_string(), "10.0.0.5");
assert_eq!(s.to_string(), "9.9.9.9");
}
}