use bytes::Bytes;
use crate::codec::{Error, header::Header, name::Name, reader::Reader};
pub const OPT_TYPE: u16 = 41;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TtlScan {
pub min_ttl: Option<u32>,
pub ttl_offsets: Vec<usize>,
}
impl TtlScan {
pub fn scan(message: &Bytes) -> Result<Self, Error> {
let mut reader = Reader::new(message.clone());
let header = Header::read(&mut reader)?;
for _ in 0..header.qdcount {
Name::skip_rr(&mut reader)?;
reader.read_u16()?; reader.read_u16()?; }
let rr_count = header.ancount as usize + header.nscount as usize + header.arcount as usize;
let mut min_ttl: Option<u32> = None;
let mut ttl_offsets: Vec<usize> = Vec::with_capacity(rr_count.min(64));
for _ in 0..rr_count {
Name::skip_rr(&mut reader)?;
let rr_type = reader.read_u16()?;
let _class = reader.read_u16()?;
let ttl_offset = reader.position();
let ttl = reader.read_u32()?;
let rdlength = reader.read_u16()? as usize;
reader.read_slice(rdlength)?;
if rr_type == OPT_TYPE {
continue;
}
ttl_offsets.push(ttl_offset);
min_ttl = Some(match min_ttl {
None => ttl,
Some(prev) => prev.min(ttl),
});
}
Ok(Self {
min_ttl,
ttl_offsets,
})
}
}
#[cfg(test)]
mod tests {
use bytes::{Bytes, BytesMut};
use super::*;
use crate::codec::{
header::{Header, Rcode},
name::Name,
reader::Reader,
writer::Writer,
};
fn write_name(w: &mut Writer, name: &str) {
let n: Name = name.parse().expect("test helper: valid name");
n.write(w);
}
fn write_ptr(w: &mut Writer, target: u16) {
w.write_u8(0xC0 | ((target >> 8) as u8));
w.write_u8((target & 0xFF) as u8);
}
fn write_a_record(w: &mut Writer, owner: &str, ttl: u32, rdata_byte: u8) -> usize {
write_name(w, owner);
w.write_u16(1); w.write_u16(1); let ttl_offset = w.len(); w.write_u32(ttl);
w.write_u16(4); w.write_slice(&[rdata_byte; 4]);
ttl_offset
}
fn write_soa_record(w: &mut Writer, owner: &str, ttl: u32) -> usize {
write_name(w, owner);
w.write_u16(6); w.write_u16(1); let ttl_offset = w.len();
w.write_u32(ttl);
let mut rdata = Vec::new();
rdata.push(0x00); rdata.push(0x00); rdata.extend_from_slice(&[0, 0, 0, 1]); rdata.extend_from_slice(&[0, 0, 0, 2]); rdata.extend_from_slice(&[0, 0, 0, 3]); rdata.extend_from_slice(&[0, 0, 0, 4]); rdata.extend_from_slice(&[0, 0, 0, 5]); w.write_u16(rdata.len() as u16);
w.write_slice(&rdata);
ttl_offset
}
fn write_opt_record(w: &mut Writer, udp_payload_size: u16) {
w.write_u8(0x00); w.write_u16(OPT_TYPE); w.write_u16(udp_payload_size); w.write_u32(0); w.write_u16(0); }
fn build_response<F>(
ancount: u16,
nscount: u16,
arcount: u16,
question_name: &str,
build_rrs: F,
) -> Bytes
where
F: FnOnce(&mut Writer),
{
let mut w = Writer::with_capacity(256);
Header::new(0xABCD)
.with_qr(true)
.with_rcode(Rcode::NoError)
.with_qdcount(1)
.with_ancount(ancount)
.with_nscount(nscount)
.with_arcount(arcount)
.write(&mut w);
write_name(&mut w, question_name);
w.write_u16(1); w.write_u16(1); build_rrs(&mut w);
w.finish()
}
fn read_u32_at(msg: &Bytes, offset: usize) -> u32 {
let mut r = Reader::new(msg.clone());
r.read_slice(offset).unwrap();
r.read_u32().unwrap()
}
#[test]
fn multi_rr_min_ttl_and_offset_count() {
let msg = build_response(3, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", 300, 0x01);
write_a_record(w, "example.com", 60, 0x02);
write_a_record(w, "example.com", 180, 0x03);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(60), "min_ttl should be 60");
assert_eq!(scan.ttl_offsets.len(), 3, "should have 3 TTL offsets");
}
#[test]
fn offsets_point_at_correct_ttl_bytes() {
let ttls = [300u32, 60, 180];
let msg = build_response(3, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", ttls[0], 0x01);
write_a_record(w, "example.com", ttls[1], 0x02);
write_a_record(w, "example.com", ttls[2], 0x03);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.ttl_offsets.len(), 3);
for (i, (&offset, &expected_ttl)) in scan.ttl_offsets.iter().zip(ttls.iter()).enumerate() {
let actual = read_u32_at(&msg, offset);
assert_eq!(
actual, expected_ttl,
"offset[{i}]={offset}: expected TTL {expected_ttl}, got {actual}"
);
}
}
#[test]
fn answer_and_authority_sections_scanned() {
let msg = build_response(2, 1, 0, "example.com", |w| {
write_a_record(w, "example.com", 120, 0x01);
write_a_record(w, "example.com", 240, 0x02);
write_soa_record(w, "example.com", 30);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(30), "SOA TTL should be the minimum");
assert_eq!(scan.ttl_offsets.len(), 3);
let expected_ttls = [120u32, 240, 30];
for (&offset, &exp) in scan.ttl_offsets.iter().zip(expected_ttls.iter()) {
assert_eq!(read_u32_at(&msg, offset), exp);
}
}
#[test]
fn patch_decrements_ttls_in_place() {
let ttls = [300u32, 60, 180];
let msg = build_response(3, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", ttls[0], 0x01);
write_a_record(w, "example.com", ttls[1], 0x02);
write_a_record(w, "example.com", ttls[2], 0x03);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
let mut patched: BytesMut = BytesMut::from(msg.as_ref());
const ELAPSED: u32 = 30;
for &offset in &scan.ttl_offsets {
let old = u32::from_be_bytes([
patched[offset],
patched[offset + 1],
patched[offset + 2],
patched[offset + 3],
]);
let new_ttl = old.saturating_sub(ELAPSED);
let bytes = new_ttl.to_be_bytes();
patched[offset] = bytes[0];
patched[offset + 1] = bytes[1];
patched[offset + 2] = bytes[2];
patched[offset + 3] = bytes[3];
}
let patched_bytes = patched.freeze();
let patched_scan = TtlScan::scan(&patched_bytes).expect("patched scan must succeed");
for (&offset, &original_ttl) in scan.ttl_offsets.iter().zip(ttls.iter()) {
let expected = original_ttl.saturating_sub(ELAPSED);
let actual = read_u32_at(&patched_bytes, offset);
assert_eq!(
actual, expected,
"offset {offset}: expected {expected}, got {actual}"
);
}
assert_eq!(patched_scan.min_ttl, Some(30));
assert_eq!(patched_scan.ttl_offsets.len(), 3);
}
#[test]
fn opt_only_response_gives_none_min_ttl() {
let msg = build_response(0, 0, 1, "example.com", |w| {
write_opt_record(w, 4096);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, None, "OPT-only: min_ttl must be None");
assert!(
scan.ttl_offsets.is_empty(),
"OPT-only: ttl_offsets must be empty"
);
}
#[test]
fn zero_rr_sections_gives_none_min_ttl() {
let msg = build_response(0, 0, 0, "example.com", |_| {});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, None);
assert!(scan.ttl_offsets.is_empty());
}
#[test]
fn a_record_plus_opt_excludes_opt_from_offsets() {
let a_ttl = 120u32;
let msg = build_response(1, 0, 1, "example.com", |w| {
write_a_record(w, "example.com", a_ttl, 0x7F);
write_opt_record(w, 4096);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(
scan.min_ttl,
Some(a_ttl),
"min_ttl should equal A record TTL"
);
assert_eq!(
scan.ttl_offsets.len(),
1,
"exactly one offset (the A record)"
);
let actual = read_u32_at(&msg, scan.ttl_offsets[0]);
assert_eq!(actual, a_ttl);
}
#[test]
fn compression_pointer_owner_name_ttl_offset_correct() {
let mut w = Writer::with_capacity(128);
Header::new(0x1234)
.with_qr(true)
.with_qdcount(1)
.with_ancount(1)
.write(&mut w);
write_name(&mut w, "example.com");
w.write_u16(1); w.write_u16(1); let question_name_offset: u16 = 12;
write_ptr(&mut w, question_name_offset);
w.write_u16(1); w.write_u16(1); let ttl_offset = w.len(); let expected_ttl: u32 = 300;
w.write_u32(expected_ttl);
w.write_u16(4); w.write_slice(&[1, 2, 3, 4]); let msg = w.finish();
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(expected_ttl));
assert_eq!(scan.ttl_offsets.len(), 1);
assert_eq!(
scan.ttl_offsets[0], ttl_offset,
"TTL offset must match the manually computed offset"
);
assert_eq!(read_u32_at(&msg, scan.ttl_offsets[0]), expected_ttl);
}
#[test]
fn pointer_loop_in_owner_name_returns_error() {
let mut w = Writer::with_capacity(128);
Header::new(0xBEEF)
.with_qr(true)
.with_qdcount(1)
.with_ancount(1)
.write(&mut w);
write_name(&mut w, "example.com");
w.write_u16(1);
w.write_u16(1);
let self_offset = w.len() as u16;
write_ptr(&mut w, self_offset);
w.write_u16(1);
w.write_u16(1);
w.write_u32(300);
w.write_u16(0);
let msg = w.finish();
let result = TtlScan::scan(&msg);
assert!(
result.is_err(),
"pointer loop / forward pointer must return an error"
);
}
#[test]
fn truncated_rdata_returns_error() {
let mut w = Writer::with_capacity(64);
Header::new(0x1234)
.with_qr(true)
.with_qdcount(1)
.with_ancount(1)
.write(&mut w);
write_name(&mut w, "example.com");
w.write_u16(1); w.write_u16(1); write_name(&mut w, "example.com");
w.write_u16(1); w.write_u16(1); w.write_u32(300); w.write_u16(100); w.write_slice(&[1, 2, 3, 4]); let msg = w.finish();
let result = TtlScan::scan(&msg);
assert!(
result.is_err(),
"truncated RDATA must return an error, not silently succeed"
);
}
#[test]
fn truncated_rr_fixed_fields_returns_error() {
let mut w = Writer::with_capacity(64);
Header::new(0x1234)
.with_qr(true)
.with_qdcount(1)
.with_ancount(1)
.write(&mut w);
write_name(&mut w, "example.com");
w.write_u16(1); w.write_u16(1); write_name(&mut w, "example.com");
w.write_u16(1); let msg = w.finish();
let result = TtlScan::scan(&msg);
assert!(
result.is_err(),
"truncated RR fixed fields must return an error"
);
}
#[test]
fn no_panic_on_random_bytes() {
let data: Vec<u8> = (0u8..=255).cycle().take(512).collect();
let msg = Bytes::from(data);
let _ = TtlScan::scan(&msg); }
#[test]
fn no_panic_on_all_zeros() {
let msg = Bytes::from(vec![0u8; 64]);
let _ = TtlScan::scan(&msg);
}
#[test]
fn no_panic_on_all_ones() {
let msg = Bytes::from(vec![0xFFu8; 256]);
let _ = TtlScan::scan(&msg);
}
#[test]
fn no_panic_on_empty() {
let msg = Bytes::new();
let result = TtlScan::scan(&msg);
assert!(
matches!(result, Err(Error::MessageTooShort(_))),
"empty buffer must return MessageTooShort"
);
}
#[test]
fn single_rr_min_ttl_equals_that_ttl() {
let msg = build_response(1, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", 42, 0x10);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(42));
assert_eq!(scan.ttl_offsets.len(), 1);
assert_eq!(read_u32_at(&msg, scan.ttl_offsets[0]), 42);
}
#[test]
fn opt_among_real_rrs_excluded_from_min_and_offsets() {
let msg = build_response(2, 0, 1, "example.com", |w| {
write_a_record(w, "example.com", 500, 0x01);
write_a_record(w, "example.com", 200, 0x02);
write_opt_record(w, 4096);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(200));
assert_eq!(scan.ttl_offsets.len(), 2);
let expected = [500u32, 200];
for (&off, &exp) in scan.ttl_offsets.iter().zip(expected.iter()) {
assert_eq!(read_u32_at(&msg, off), exp);
}
}
#[test]
fn ttl_zero_is_recorded_and_is_minimum() {
let msg = build_response(2, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", 300, 0x01);
write_a_record(w, "example.com", 0, 0x02);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(0));
assert_eq!(scan.ttl_offsets.len(), 2);
}
#[test]
fn ttl_max_u32_recorded() {
let msg = build_response(1, 0, 0, "example.com", |w| {
write_a_record(w, "example.com", u32::MAX, 0x01);
});
let scan = TtlScan::scan(&msg).expect("scan must succeed");
assert_eq!(scan.min_ttl, Some(u32::MAX));
assert_eq!(read_u32_at(&msg, scan.ttl_offsets[0]), u32::MAX);
}
}