103 lines
2.9 KiB
Rust
103 lines
2.9 KiB
Rust
use dhcproto::Decodable;
|
|
use dhcproto::v4::{DhcpOption, MessageType, OptionCode};
|
|
use smoltcp::wire::Ipv4Address;
|
|
use std::collections::HashSet;
|
|
use std::time::Duration;
|
|
|
|
#[derive(Default)]
|
|
pub struct DhcpSnooper {
|
|
vm_lease: Option<Lease>,
|
|
uncertainty_duration: Duration,
|
|
}
|
|
|
|
impl DhcpSnooper {
|
|
pub fn new(uncertainty_duration: Duration) -> Self {
|
|
DhcpSnooper {
|
|
uncertainty_duration,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
pub fn register_dhcp_reply(&mut self, dhcp_packet: &[u8]) {
|
|
let mut decoder = dhcproto::v4::Decoder::new(dhcp_packet);
|
|
|
|
let message = match dhcproto::v4::Message::decode(&mut decoder) {
|
|
Ok(message) => message,
|
|
Err(_) => return,
|
|
};
|
|
|
|
match message.opts().msg_type() {
|
|
Some(MessageType::Ack) => {
|
|
let lease_time = match message.opts().get(OptionCode::AddressLeaseTime) {
|
|
Some(DhcpOption::AddressLeaseTime(lease_time)) => lease_time,
|
|
_ => return,
|
|
};
|
|
|
|
let dns_ips = match message.opts().get(OptionCode::DomainNameServer) {
|
|
Some(DhcpOption::DomainNameServer(dns_ips)) => {
|
|
HashSet::from_iter(dns_ips.iter().cloned())
|
|
}
|
|
_ => HashSet::new(),
|
|
};
|
|
|
|
let mut lease_duration = Duration::from_secs(*lease_time as u64);
|
|
|
|
// Adjust for uncertainty caused by using a coarse clock
|
|
lease_duration = lease_duration.saturating_sub(self.uncertainty_duration);
|
|
|
|
self.vm_lease = Some(Lease::new(message.yiaddr(), lease_duration, dns_ips))
|
|
}
|
|
Some(MessageType::Nak) => {
|
|
self.vm_lease = None;
|
|
}
|
|
_ => {}
|
|
};
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) fn set_lease(&mut self, vm_lease: Option<Lease>) {
|
|
self.vm_lease = vm_lease
|
|
}
|
|
|
|
pub fn lease(&self) -> &Option<Lease> {
|
|
&self.vm_lease
|
|
}
|
|
|
|
pub fn valid_dns_target(&self, addr: &Ipv4Address) -> bool {
|
|
if let Some(lease) = &self.vm_lease {
|
|
return lease.dns_ips.contains(addr);
|
|
}
|
|
|
|
false
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct Lease {
|
|
address: Ipv4Address,
|
|
valid_until: coarsetime::Instant,
|
|
dns_ips: HashSet<Ipv4Address>,
|
|
}
|
|
|
|
impl Lease {
|
|
pub fn new(address: Ipv4Address, lease_time: Duration, dns_ips: HashSet<Ipv4Address>) -> Lease {
|
|
Lease {
|
|
address,
|
|
valid_until: coarsetime::Instant::recent() + lease_time.into(),
|
|
dns_ips,
|
|
}
|
|
}
|
|
|
|
pub fn address(&self) -> Ipv4Address {
|
|
self.address
|
|
}
|
|
|
|
pub fn valid(&self) -> bool {
|
|
coarsetime::Instant::recent() < self.valid_until
|
|
}
|
|
|
|
pub fn valid_ip_source(&self, address: Ipv4Address) -> bool {
|
|
self.address == address && self.valid()
|
|
}
|
|
}
|