Skip to main content

stackforge_core/layer/dns/
bitmap.rs

1//! NSEC/NSEC3 type bitmap encoding and decoding (RFC 4034 Section 4.1.2).
2
3use crate::layer::field::FieldError;
4
5/// Convert NSEC/NSEC3 type bitmap wire format to a list of RR type numbers.
6///
7/// The bitmap format consists of windows:
8///   - Window block number (1 byte): high 8 bits of type number
9///   - Bitmap length (1 byte): number of bitmap bytes (1-32)
10///   - Bitmap (variable): each bit represents a type in that window
11pub fn bitmap_to_rr_list(data: &[u8]) -> Result<Vec<u16>, FieldError> {
12    let mut types = Vec::new();
13    let mut pos = 0;
14
15    while pos < data.len() {
16        if pos + 2 > data.len() {
17            return Err(FieldError::InvalidValue(
18                "truncated NSEC bitmap window".to_string(),
19            ));
20        }
21
22        let window = u16::from(data[pos]);
23        let bitmap_len = data[pos + 1] as usize;
24        pos += 2;
25
26        if bitmap_len == 0 || bitmap_len > 32 {
27            return Err(FieldError::InvalidValue(format!(
28                "invalid NSEC bitmap length: {bitmap_len}"
29            )));
30        }
31
32        if pos + bitmap_len > data.len() {
33            return Err(FieldError::InvalidValue(
34                "truncated NSEC bitmap data".to_string(),
35            ));
36        }
37
38        for i in 0..bitmap_len {
39            let byte = data[pos + i];
40            for bit in 0..8u16 {
41                if (byte >> (7 - bit)) & 1 != 0 {
42                    let rr_type = window * 256 + (i as u16) * 8 + bit;
43                    types.push(rr_type);
44                }
45            }
46        }
47
48        pos += bitmap_len;
49    }
50
51    Ok(types)
52}
53
54/// Convert a list of RR type numbers to NSEC/NSEC3 type bitmap wire format.
55#[must_use]
56pub fn rr_list_to_bitmap(types: &[u16]) -> Vec<u8> {
57    if types.is_empty() {
58        return Vec::new();
59    }
60
61    // Group types by window (high byte)
62    let mut windows: std::collections::BTreeMap<u8, Vec<u16>> = std::collections::BTreeMap::new();
63    for &t in types {
64        let window = (t >> 8) as u8;
65        let offset = t & 0xFF;
66        windows.entry(window).or_default().push(offset);
67    }
68
69    let mut out = Vec::new();
70
71    for (&window, offsets) in &windows {
72        // Find the maximum byte needed in this window
73        let max_offset = *offsets.iter().max().unwrap();
74        let bitmap_len = (max_offset / 8 + 1) as usize;
75
76        // Build the bitmap
77        let mut bitmap = vec![0u8; bitmap_len];
78        for &offset in offsets {
79            let byte_idx = (offset / 8) as usize;
80            let bit_idx = 7 - (offset % 8);
81            bitmap[byte_idx] |= 1 << bit_idx;
82        }
83
84        out.push(window);
85        out.push(bitmap_len as u8);
86        out.extend_from_slice(&bitmap);
87    }
88
89    out
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_bitmap_roundtrip_simple() {
98        let types = vec![1, 2, 5, 6, 15, 16, 28]; // A, NS, CNAME, SOA, MX, TXT, AAAA
99        let bitmap = rr_list_to_bitmap(&types);
100        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
101        assert_eq!(decoded, types);
102    }
103
104    #[test]
105    fn test_bitmap_single_type() {
106        let types = vec![1]; // Just A
107        let bitmap = rr_list_to_bitmap(&types);
108        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
109        assert_eq!(decoded, types);
110    }
111
112    #[test]
113    fn test_bitmap_type_zero() {
114        let types = vec![0];
115        let bitmap = rr_list_to_bitmap(&types);
116        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
117        assert_eq!(decoded, types);
118    }
119
120    #[test]
121    fn test_bitmap_high_types() {
122        // Types in higher windows (256, 512, etc.)
123        let types = vec![256, 512, 4096, 36864];
124        let bitmap = rr_list_to_bitmap(&types);
125        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
126        assert_eq!(decoded, types);
127    }
128
129    #[test]
130    fn test_bitmap_type_65535() {
131        let types = vec![65535];
132        let bitmap = rr_list_to_bitmap(&types);
133        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
134        assert_eq!(decoded, types);
135    }
136
137    #[test]
138    fn test_bitmap_mixed_windows() {
139        let mut types = vec![1, 2, 28, 46, 47, 48, 256, 257];
140        let bitmap = rr_list_to_bitmap(&types);
141        let decoded = bitmap_to_rr_list(&bitmap).unwrap();
142        types.sort();
143        assert_eq!(decoded, types);
144    }
145
146    #[test]
147    fn test_bitmap_empty() {
148        let bitmap = rr_list_to_bitmap(&[]);
149        assert!(bitmap.is_empty());
150    }
151
152    #[test]
153    fn test_bitmap_decode_truncated() {
154        let data = vec![0]; // Only window byte, no length
155        assert!(bitmap_to_rr_list(&data).is_err());
156    }
157
158    #[test]
159    fn test_bitmap_decode_invalid_length() {
160        let data = vec![0, 33]; // Length > 32
161        assert!(bitmap_to_rr_list(&data).is_err());
162    }
163
164    #[test]
165    fn test_bitmap_decode_zero_length() {
166        let data = vec![0, 0]; // Length == 0
167        assert!(bitmap_to_rr_list(&data).is_err());
168    }
169}