simple_dns/dns/rdata/
txt.rs

1use crate::CharacterString;
2use crate::{
3    dns::{WireFormat, MAX_CHARACTER_STRING_LENGTH},
4    lib::{vec, FromUtf8Error, String, Vec},
5    lib::Write,
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            let value = splited.next();
69            Some((key, value))
70        })
71    }
72
73    // FIXME: remove the std feature once the HashMap is sorted out
74    /// Returns parsed attributes from this TXT Record, valid formats are:
75    /// - key=value
76    /// - key=
77    /// - key
78    ///
79    /// If a key is duplicated, only the first one will be considered
80    #[cfg(feature = "std")]
81    pub fn attributes(&self) -> crate::lib::HashMap<String, Option<String>> {
82        let mut attributes = crate::lib::HashMap::new();
83        let iter = self.iter_raw().filter_map(|(key, value)| {
84            let key = match crate::lib::str::from_utf8(key) {
85                Ok(key) => key.to_owned(),
86                Err(_) => return None,
87            };
88
89            let value = match value {
90                Some(value) if !value.is_empty() => match crate::lib::str::from_utf8(value) {
91                    Ok(v) => Some(v.to_owned()),
92                    Err(_) => Some(String::new()),
93                },
94                Some(_) => Some(String::new()),
95                _ => None,
96            };
97
98            Some((key, value))
99        });
100
101        for (key, value) in iter {
102            attributes.entry(key).or_insert(value);
103        }
104
105        attributes
106    }
107
108    /// Similar to [`attributes()`](TXT::attributes) but it parses the full TXT record as a single string,
109    /// instead of expecting each attribute to be a separate [`CharacterString`](`CharacterString`)
110    #[cfg(feature = "std")]
111    pub fn long_attributes(self) -> crate::Result<crate::lib::HashMap<String, Option<String>>> {
112        let mut attributes = crate::lib::HashMap::new();
113
114        let full_string: String = match self.try_into() {
115            Ok(string) => string,
116            Err(err) => return Err(crate::SimpleDnsError::InvalidUtf8String(err)),
117        };
118
119        let parts = full_string.split(|c| (c as u8) == b';');
120
121        for part in parts {
122            let key_value = part.splitn(2, |c| (c as u8) == b'=').collect::<Vec<&str>>();
123
124            let key = key_value[0];
125
126            let value = match key_value.len() > 1 {
127                true => Some(key_value[1].to_owned()),
128                _ => None,
129            };
130
131            if !key.is_empty() {
132                attributes.entry(key.to_owned()).or_insert(value);
133            }
134        }
135
136        Ok(attributes)
137    }
138
139    /// Transforms the inner data into its owned type
140    pub fn into_owned<'b>(self) -> TXT<'b> {
141        TXT {
142            strings: self.strings.into_iter().map(|s| s.into_owned()).collect(),
143            size: self.size,
144        }
145    }
146}
147
148#[cfg(feature = "std")]
149impl TryFrom<crate::lib::HashMap<String, Option<String>>> for TXT<'_> {
150    type Error = crate::SimpleDnsError;
151
152    fn try_from(value: crate::lib::HashMap<String, Option<String>>) -> Result<Self, Self::Error> {
153        let mut txt = TXT::new();
154        for (key, value) in value {
155            match value {
156                Some(value) => {
157                    txt.add_char_string(format!("{}={}", &key, &value).try_into()?);
158                }
159                None => txt.add_char_string(key.try_into()?),
160            }
161        }
162        Ok(txt)
163    }
164}
165
166impl<'a> TryFrom<&'a str> for TXT<'a> {
167    type Error = crate::SimpleDnsError;
168
169    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
170        let mut txt = TXT::new();
171        for v in value.as_bytes().chunks(MAX_CHARACTER_STRING_LENGTH - 1) {
172            txt.add_char_string(CharacterString::new(v)?);
173        }
174        Ok(txt)
175    }
176}
177
178impl<'a> TryFrom<TXT<'a>> for String {
179    type Error = FromUtf8Error;
180
181    fn try_from(val: TXT<'a>) -> Result<Self, Self::Error> {
182        let init = Vec::with_capacity(val.len());
183
184        let bytes = val.strings.into_iter().fold(init, |mut acc, val| {
185            acc.extend(val.data.as_ref());
186            acc
187        });
188        String::from_utf8(bytes)
189    }
190}
191
192impl<'a> WireFormat<'a> for TXT<'a> {
193    const MINIMUM_LEN: usize = 1;
194
195    fn parse(data: &mut crate::bytes_buffer::BytesBuffer<'a>) -> crate::Result<Self>
196    where
197        Self: Sized,
198    {
199        let mut strings = Vec::new();
200        let mut size = 0;
201
202        while data.has_remaining() {
203            let char_str = CharacterString::parse(data)?;
204            size += char_str.len();
205            strings.push(char_str);
206        }
207
208        Ok(Self { strings, size })
209    }
210
211    fn len(&self) -> usize {
212        if self.strings.is_empty() {
213            Self::MINIMUM_LEN
214        } else {
215            self.size
216        }
217    }
218
219    fn write_to<T: Write>(&self, out: &mut T) -> crate::Result<()> {
220        if self.strings.is_empty() {
221            out.write_all(&[0])?;
222        } else {
223            for string in &self.strings {
224                string.write_to(out)?;
225            }
226        }
227        Ok(())
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use crate::lib::Error;
235
236    #[test]
237    pub fn parse_and_write_txt() -> Result<(), crate::lib::Box<dyn Error>> {
238        let mut out = vec![];
239        let txt = TXT::new()
240            .with_char_string("version=0.1".try_into()?)
241            .with_char_string("proto=123".try_into()?);
242
243        txt.write_to(&mut out)?;
244        assert_eq!(out.len(), txt.len());
245
246        let txt2 = TXT::parse(&mut out[..].into())?;
247        assert_eq!(2, txt2.strings.len());
248        assert_eq!(txt.strings[0], txt2.strings[0]);
249        assert_eq!(txt.strings[1], txt2.strings[1]);
250
251        Ok(())
252    }
253
254    #[test]
255    pub fn iter_raw() -> Result<(), crate::lib::Box<dyn Error>> {
256        let txt = TXT::new()
257            .with_string("version=0.1")?
258            .with_string("flag")?
259            .with_string("with_eq=eq=")?
260            .with_string("version=dup")?
261            .with_string("empty=")?;
262
263        assert_eq!(
264            txt.iter_raw().collect::<Vec<_>>(),
265            vec![
266                ("version".as_bytes(), Some("0.1".as_bytes())),
267                ("flag".as_bytes(), None),
268                ("with_eq".as_bytes(), Some("eq=".as_bytes())),
269                ("version".as_bytes(), Some("dup".as_bytes())),
270                ("empty".as_bytes(), Some("".as_bytes()))
271            ]
272        );
273        Ok(())
274    }
275
276    #[test]
277    #[cfg(feature = "std")]
278    pub fn get_attributes() -> Result<(), Box<dyn Error>> {
279        let attributes = TXT::new()
280            .with_string("version=0.1")?
281            .with_string("flag")?
282            .with_string("with_eq=eq=")?
283            .with_string("version=dup")?
284            .with_string("empty=")?
285            .attributes();
286
287        assert_eq!(4, attributes.len());
288        assert_eq!(Some("0.1".to_owned()), attributes["version"]);
289        assert_eq!(Some("eq=".to_owned()), attributes["with_eq"]);
290        assert_eq!(Some(String::new()), attributes["empty"]);
291        assert_eq!(None, attributes["flag"]);
292
293        Ok(())
294    }
295
296    #[test]
297    #[cfg(feature = "std")]
298    fn parse_sample() -> Result<(), Box<dyn std::error::Error>> {
299        use crate::{rdata::RData, ResourceRecord};
300        let sample_file = std::fs::read("samples/zonefile/TXT.sample")?;
301
302        let sample_rdata = match ResourceRecord::parse(&mut sample_file[..].into())?.rdata {
303            RData::TXT(rdata) => rdata,
304            _ => unreachable!(),
305        };
306
307        let strings = vec!["\"foo\nbar\"".try_into()?];
308        assert_eq!(sample_rdata.strings, strings);
309
310        Ok(())
311    }
312
313    #[test]
314    fn write_and_parse_large_txt() -> Result<(), crate::lib::Box<dyn Error>> {
315        let string = "X".repeat(1000);
316        let txt: TXT = string.as_str().try_into()?;
317
318        let mut bytes = Vec::new();
319        assert!(txt.write_to(&mut bytes).is_ok());
320
321        let parsed_txt = TXT::parse(&mut bytes[..].into())?;
322        let parsed_string: String = parsed_txt.try_into()?;
323
324        assert_eq!(parsed_string, string);
325
326        Ok(())
327    }
328
329    #[test]
330    #[cfg(feature = "std")]
331    fn write_and_parse_large_attributes() -> Result<(), Box<dyn Error>> {
332        let big_value = "f".repeat(1000);
333
334        let string = format!("foo={big_value};;flag;bar={big_value}");
335        let txt: TXT = string.as_str().try_into()?;
336        let attributes = txt.long_attributes()?;
337
338        assert_eq!(Some(big_value.to_owned()), attributes["bar"]);
339
340        Ok(())
341    }
342}
simple_dns/dns/rdata/txt.rs

simple_dns/dns/rdata/
txt.rs
