simple_dns/dns/rdata/
txt.rs

1use crate::CharacterString;
2use crate::{
3    dns::{WireFormat, MAX_CHARACTER_STRING_LENGTH},
4    lib::Write,
5    lib::{vec, FromUtf8Error, String, Vec},
6};
7
8use super::RR;
9
10/// Represents a TXT Resource Record
11#[derive(Debug, PartialEq, Eq, Hash, Clone)]
12pub struct TXT<'a> {
13    strings: Vec<CharacterString<'a>>,
14    size: usize,
15}
16
17impl RR for TXT<'_> {
18    const TYPE_CODE: u16 = 16;
19}
20
21impl Default for TXT<'_> {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl<'a> TXT<'a> {
28    /// Creates a new empty TXT Record
29    pub fn new() -> Self {
30        Self {
31            strings: vec![],
32            size: 0,
33        }
34    }
35
36    /// Add `char_string` to this TXT record as a validated [`CharacterString`](`CharacterString`)
37    pub fn add_string(&mut self, char_string: &'a str) -> crate::Result<()> {
38        self.add_char_string(char_string.try_into()?);
39        Ok(())
40    }
41
42    /// Add `char_string` to this TXT record
43    pub fn add_char_string(&mut self, char_string: CharacterString<'a>) {
44        self.size += char_string.len();
45        self.strings.push(char_string);
46    }
47
48    /// Add `char_string` to this TXT record as a validated [`CharacterString`](`CharacterString`), consuming and returning Self
49    pub fn with_string(mut self, char_string: &'a str) -> crate::Result<Self> {
50        self.add_char_string(char_string.try_into()?);
51        Ok(self)
52    }
53
54    /// Add `char_string` to this TXT record, consuming and returning Self
55    pub fn with_char_string(mut self, char_string: CharacterString<'a>) -> Self {
56        self.add_char_string(char_string);
57        self
58    }
59
60    /// Returns parsed attributes from this TXT Record as bytes, valid formats are:
61    /// - key=value
62    /// - key=
63    /// - key
64    pub fn iter_raw(&self) -> impl Iterator<Item = (&[u8], Option<&[u8]>)> {
65        self.strings.iter().filter_map(|char_str| {
66            let mut splited = char_str.data.splitn(2, |c| *c == b'=');
67            let key = splited.next()?;
68            if key.is_empty() {
69                return None;
70            }
71            let value = splited.next();
72            Some((key, value))
73        })
74    }
75
76    /// Add a key/value attribute to this TXT record per RFC 6763 Section 6.
77    ///
78    /// The `key` must be non-empty printable US-ASCII (`0x20..=0x7E`) excluding `=`.
79    /// The optional `value` may contain arbitrary binary bytes, as long as the combined
80    /// `key=value` fits within a 255-byte [`CharacterString`].
81    pub fn add_attribute(&mut self, key: &str, value: Option<&[u8]>) -> crate::Result<()> {
82        if key.is_empty()
83            || !key
84                .bytes()
85                .all(|b| (0x20..=0x7E).contains(&b) && b != b'=')
86        {
87            return Err(crate::SimpleDnsError::InvalidCharacterString);
88        }
89
90        let total_len = key.len() + value.map_or(0, |v| 1 + v.len());
91        if total_len > MAX_CHARACTER_STRING_LENGTH {
92            return Err(crate::SimpleDnsError::InvalidCharacterString);
93        }
94
95        let mut data = Vec::with_capacity(total_len);
96        data.extend_from_slice(key.as_bytes());
97        if let Some(value) = value {
98            data.push(b'=');
99            data.extend_from_slice(value);
100        }
101
102        self.add_char_string(CharacterString::new(&data)?.into_owned());
103        Ok(())
104    }
105
106    /// Add a key/value attribute to this TXT record, consuming and returning `Self`.
107    pub fn with_attribute(mut self, key: &str, value: Option<&[u8]>) -> crate::Result<Self> {
108        self.add_attribute(key, value)?;
109        Ok(self)
110    }
111
112    // FIXME: remove the std feature once the HashMap is sorted out
113    /// Returns parsed attributes from this TXT Record, valid formats are:
114    /// - key=value
115    /// - key=
116    /// - key
117    ///
118    /// If a key is duplicated, only the first one will be considered (case-insensitive per RFC 6763 Section 6.4)
119    #[cfg(feature = "std")]
120    pub fn attributes(&self) -> crate::lib::HashMap<String, Option<String>> {
121        let mut attributes = crate::lib::HashMap::new();
122        let iter = self.iter_raw().filter_map(|(key, value)| {
123            let key = match crate::lib::str::from_utf8(key) {
124                Ok(key) => key.to_ascii_lowercase(),
125                Err(_) => return None,
126            };
127
128            let value = match value {
129                Some(value) if !value.is_empty() => match crate::lib::str::from_utf8(value) {
130                    Ok(v) => Some(v.to_owned()),
131                    Err(_) => Some(String::new()),
132                },
133                Some(_) => Some(String::new()),
134                _ => None,
135            };
136
137            Some((key, value))
138        });
139
140        for (key, value) in iter {
141            attributes.entry(key).or_insert(value);
142        }
143
144        attributes
145    }
146
147    /// Similar to [`attributes()`](TXT::attributes) but it parses the full TXT record as a single string,
148    /// instead of expecting each attribute to be a separate [`CharacterString`](`CharacterString`)
149    #[cfg(feature = "std")]
150    pub fn long_attributes(self) -> crate::Result<crate::lib::HashMap<String, Option<String>>> {
151        let mut attributes = crate::lib::HashMap::new();
152
153        let full_string: String = match self.try_into() {
154            Ok(string) => string,
155            Err(err) => return Err(crate::SimpleDnsError::InvalidUtf8String(err)),
156        };
157
158        let parts = full_string.split(|c| (c as u8) == b';');
159
160        for part in parts {
161            let key_value = part.splitn(2, |c| (c as u8) == b'=').collect::<Vec<&str>>();
162
163            let key = key_value[0];
164
165            let value = match key_value.len() > 1 {
166                true => Some(key_value[1].to_owned()),
167                _ => None,
168            };
169
170            if !key.is_empty() {
171                attributes.entry(key.to_owned()).or_insert(value);
172            }
173        }
174
175        Ok(attributes)
176    }
177
178    /// Transforms the inner data into its owned type
179    pub fn into_owned<'b>(self) -> TXT<'b> {
180        TXT {
181            strings: self.strings.into_iter().map(|s| s.into_owned()).collect(),
182            size: self.size,
183        }
184    }
185}
186
187#[cfg(feature = "std")]
188impl TryFrom<crate::lib::HashMap<String, Option<String>>> for TXT<'_> {
189    type Error = crate::SimpleDnsError;
190
191    fn try_from(value: crate::lib::HashMap<String, Option<String>>) -> Result<Self, Self::Error> {
192        let mut txt = TXT::new();
193        for (key, value) in value {
194            match value {
195                Some(value) => {
196                    txt.add_char_string(format!("{key}={value}").try_into()?);
197                }
198                None => txt.add_char_string(key.try_into()?),
199            }
200        }
201        Ok(txt)
202    }
203}
204
205impl<'a> TryFrom<&'a str> for TXT<'a> {
206    type Error = crate::SimpleDnsError;
207
208    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
209        let mut txt = TXT::new();
210        for v in value.as_bytes().chunks(MAX_CHARACTER_STRING_LENGTH - 1) {
211            txt.add_char_string(CharacterString::new(v)?);
212        }
213        Ok(txt)
214    }
215}
216
217impl<'a> TryFrom<TXT<'a>> for String {
218    type Error = FromUtf8Error;
219
220    fn try_from(val: TXT<'a>) -> Result<Self, Self::Error> {
221        let init = Vec::with_capacity(val.len());
222
223        let bytes = val.strings.into_iter().fold(init, |mut acc, val| {
224            acc.extend(val.data.as_ref());
225            acc
226        });
227        String::from_utf8(bytes)
228    }
229}
230
231impl<'a> WireFormat<'a> for TXT<'a> {
232    const MINIMUM_LEN: usize = 1;
233
234    fn parse(data: &mut crate::bytes_buffer::BytesBuffer<'a>) -> crate::Result<Self>
235    where
236        Self: Sized,
237    {
238        let mut strings = Vec::new();
239        let mut size = 0;
240
241        while data.has_remaining() {
242            let char_str = CharacterString::parse(data)?;
243            size += char_str.len();
244            strings.push(char_str);
245        }
246
247        Ok(Self { strings, size })
248    }
249
250    fn len(&self) -> usize {
251        if self.strings.is_empty() {
252            Self::MINIMUM_LEN
253        } else {
254            self.size
255        }
256    }
257
258    fn write_to<T: Write>(&self, out: &mut T) -> crate::Result<()> {
259        if self.strings.is_empty() {
260            out.write_all(&[0])?;
261        } else {
262            for string in &self.strings {
263                string.write_to(out)?;
264            }
265        }
266        Ok(())
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273    use crate::lib::Error;
274
275    #[test]
276    pub fn parse_and_write_txt() -> Result<(), crate::lib::Box<dyn Error>> {
277        let mut out = vec![];
278        let txt = TXT::new()
279            .with_char_string("version=0.1".try_into()?)
280            .with_char_string("proto=123".try_into()?);
281
282        txt.write_to(&mut out)?;
283        assert_eq!(out.len(), txt.len());
284
285        let txt2 = TXT::parse(&mut out[..].into())?;
286        assert_eq!(2, txt2.strings.len());
287        assert_eq!(txt.strings[0], txt2.strings[0]);
288        assert_eq!(txt.strings[1], txt2.strings[1]);
289
290        Ok(())
291    }
292
293    #[test]
294    pub fn iter_raw() -> Result<(), crate::lib::Box<dyn Error>> {
295        let txt = TXT::new()
296            .with_string("version=0.1")?
297            .with_string("flag")?
298            .with_string("with_eq=eq=")?
299            .with_string("version=dup")?
300            .with_string("empty=")?;
301
302        assert_eq!(
303            txt.iter_raw().collect::<Vec<_>>(),
304            vec![
305                ("version".as_bytes(), Some("0.1".as_bytes())),
306                ("flag".as_bytes(), None),
307                ("with_eq".as_bytes(), Some("eq=".as_bytes())),
308                ("version".as_bytes(), Some("dup".as_bytes())),
309                ("empty".as_bytes(), Some("".as_bytes()))
310            ]
311        );
312        Ok(())
313    }
314
315    #[test]
316    #[cfg(feature = "std")]
317    pub fn get_attributes() -> Result<(), Box<dyn Error>> {
318        let attributes = TXT::new()
319            .with_string("version=0.1")?
320            .with_string("flag")?
321            .with_string("with_eq=eq=")?
322            .with_string("version=dup")?
323            .with_string("empty=")?
324            .attributes();
325
326        assert_eq!(4, attributes.len());
327        assert_eq!(Some("0.1".to_owned()), attributes["version"]);
328        assert_eq!(Some("eq=".to_owned()), attributes["with_eq"]);
329        assert_eq!(Some(String::new()), attributes["empty"]);
330        assert_eq!(None, attributes["flag"]);
331
332        Ok(())
333    }
334
335    #[test]
336    #[cfg(feature = "std")]
337    fn parse_sample() -> Result<(), Box<dyn std::error::Error>> {
338        use crate::{rdata::RData, ResourceRecord};
339        let sample_file = std::fs::read("samples/zonefile/TXT.sample")?;
340
341        let sample_rdata = match ResourceRecord::parse(&mut sample_file[..].into())?.rdata {
342            RData::TXT(rdata) => rdata,
343            _ => unreachable!(),
344        };
345
346        let strings = vec!["\"foo\nbar\"".try_into()?];
347        assert_eq!(sample_rdata.strings, strings);
348
349        Ok(())
350    }
351
352    #[test]
353    fn write_and_parse_large_txt() -> Result<(), crate::lib::Box<dyn Error>> {
354        let string = "X".repeat(1000);
355        let txt: TXT = string.as_str().try_into()?;
356
357        let mut bytes = Vec::new();
358        assert!(txt.write_to(&mut bytes).is_ok());
359
360        let parsed_txt = TXT::parse(&mut bytes[..].into())?;
361        let parsed_string: String = parsed_txt.try_into()?;
362
363        assert_eq!(parsed_string, string);
364
365        Ok(())
366    }
367
368    #[test]
369    #[cfg(feature = "std")]
370    fn write_and_parse_large_attributes() -> Result<(), Box<dyn Error>> {
371        let big_value = "f".repeat(1000);
372
373        let string = format!("foo={big_value};;flag;bar={big_value}");
374        let txt: TXT = string.as_str().try_into()?;
375        let attributes = txt.long_attributes()?;
376
377        assert_eq!(Some(big_value.to_owned()), attributes["bar"]);
378
379        Ok(())
380    }
381
382    #[test]
383    #[cfg(feature = "std")]
384    fn test_rfc6763_empty_keys_and_case_insensitivity() -> Result<(), Box<dyn Error>> {
385        let txt = TXT::new()
386            .with_char_string(CharacterString::new(b"").unwrap())
387            .with_char_string(CharacterString::new(b"=ignored").unwrap())
388            .with_string("Vers=1")?
389            .with_string("vers=2")?
390            .with_string("PATH=/test")?;
391
392        let raw: Vec<_> = txt.iter_raw().collect();
393        assert_eq!(3, raw.len());
394        assert_eq!(b"Vers", raw[0].0);
395        assert_eq!(Some(b"1".as_slice()), raw[0].1);
396
397        let attrs = txt.attributes();
398        assert_eq!(2, attrs.len());
399        assert_eq!(Some("1".to_owned()), attrs["vers"]);
400        assert_eq!(Some("/test".to_owned()), attrs["path"]);
401
402        Ok(())
403    }
404
405    #[test]
406    fn test_add_attribute() -> Result<(), Box<dyn Error>> {
407        let mut txt = TXT::new();
408        txt.add_attribute("binary", Some(&[0x00, 0xFF, 0xFE]))?;
409        txt.add_attribute("flag", None)?;
410
411        let raw: Vec<_> = txt.iter_raw().collect();
412        assert_eq!(2, raw.len());
413        assert_eq!(b"binary", raw[0].0);
414        assert_eq!(Some(&[0x00, 0xFF, 0xFE][..]), raw[0].1);
415        assert_eq!(b"flag", raw[1].0);
416        assert_eq!(None, raw[1].1);
417
418        assert!(txt.clone().with_attribute("", None).is_err());
419        assert!(txt.with_attribute("bad=key", None).is_err());
420
421        Ok(())
422    }
423}
simple_dns/dns/rdata/txt.rs

simple_dns/dns/rdata/
txt.rs
