simple_dns/dns/rdata/
txt.rs

1use std::{
2    collections::HashMap,
3    convert::{TryFrom, TryInto},
4};
5
6use crate::dns::{WireFormat, MAX_CHARACTER_STRING_LENGTH};
7use crate::CharacterString;
8
9use super::RR;
10
11/// Represents a TXT Resource Record
12#[derive(Debug, PartialEq, Eq, Hash, Clone)]
13pub struct TXT<'a> {
14    strings: Vec<CharacterString<'a>>,
15    size: usize,
16}
17
18impl RR for TXT<'_> {
19    const TYPE_CODE: u16 = 16;
20}
21
22impl Default for TXT<'_> {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl<'a> TXT<'a> {
29    /// Creates a new empty TXT Record
30    pub fn new() -> Self {
31        Self {
32            strings: vec![],
33            size: 0,
34        }
35    }
36
37    /// Add `char_string` to this TXT record as a validated [`CharacterString`](`CharacterString`)
38    pub fn add_string(&mut self, char_string: &'a str) -> crate::Result<()> {
39        self.add_char_string(char_string.try_into()?);
40        Ok(())
41    }
42
43    /// Add `char_string` to this TXT record
44    pub fn add_char_string(&mut self, char_string: CharacterString<'a>) {
45        self.size += char_string.len();
46        self.strings.push(char_string);
47    }
48
49    /// Add `char_string` to this TXT record as a validated [`CharacterString`](`CharacterString`), consuming and returning Self
50    pub fn with_string(mut self, char_string: &'a str) -> crate::Result<Self> {
51        self.add_char_string(char_string.try_into()?);
52        Ok(self)
53    }
54
55    /// Add `char_string` to this TXT record, consuming and returning Self
56    pub fn with_char_string(mut self, char_string: CharacterString<'a>) -> Self {
57        self.add_char_string(char_string);
58        self
59    }
60
61    /// Returns parsed attributes from this TXT Record, valid formats are:
62    /// - key=value
63    /// - key=
64    /// - key
65    ///
66    /// If a key is duplicated, only the first one will be considered
67    pub fn attributes(&self) -> HashMap<String, Option<String>> {
68        let mut attributes = HashMap::new();
69
70        for char_str in &self.strings {
71            let mut splited = char_str.data.splitn(2, |c| *c == b'=');
72            let key = match splited.next() {
73                Some(key) => match std::str::from_utf8(key) {
74                    Ok(key) => key.to_owned(),
75                    Err(_) => continue,
76                },
77                None => continue,
78            };
79
80            let value = match splited.next() {
81                Some(value) if !value.is_empty() => match std::str::from_utf8(value) {
82                    Ok(v) => Some(v.to_owned()),
83                    Err(_) => Some(String::new()),
84                },
85                Some(_) => Some(String::new()),
86                _ => None,
87            };
88
89            attributes.entry(key).or_insert(value);
90        }
91
92        attributes
93    }
94
95    /// Similar to [`attributes()`](TXT::attributes) but it parses the full TXT record as a single string,
96    /// instead of expecting each attribute to be a separate [`CharacterString`](`CharacterString`)
97    pub fn long_attributes(self) -> crate::Result<HashMap<String, Option<String>>> {
98        let mut attributes = HashMap::new();
99
100        let full_string: String = match self.try_into() {
101            Ok(string) => string,
102            Err(err) => return Err(crate::SimpleDnsError::InvalidUtf8String(err)),
103        };
104
105        let parts = full_string.split(|c| (c as u8) == b';');
106
107        for part in parts {
108            let key_value = part.splitn(2, |c| (c as u8) == b'=').collect::<Vec<&str>>();
109
110            let key = key_value[0];
111
112            let value = match key_value.len() > 1 {
113                true => Some(key_value[1].to_owned()),
114                _ => None,
115            };
116
117            if !key.is_empty() {
118                attributes.entry(key.to_owned()).or_insert(value);
119            }
120        }
121
122        Ok(attributes)
123    }
124
125    /// Transforms the inner data into its owned type
126    pub fn into_owned<'b>(self) -> TXT<'b> {
127        TXT {
128            strings: self.strings.into_iter().map(|s| s.into_owned()).collect(),
129            size: self.size,
130        }
131    }
132}
133
134impl TryFrom<HashMap<String, Option<String>>> for TXT<'_> {
135    type Error = crate::SimpleDnsError;
136
137    fn try_from(value: HashMap<String, Option<String>>) -> Result<Self, Self::Error> {
138        let mut txt = TXT::new();
139        for (key, value) in value {
140            match value {
141                Some(value) => {
142                    txt.add_char_string(format!("{}={}", &key, &value).try_into()?);
143                }
144                None => txt.add_char_string(key.try_into()?),
145            }
146        }
147        Ok(txt)
148    }
149}
150
151impl<'a> TryFrom<&'a str> for TXT<'a> {
152    type Error = crate::SimpleDnsError;
153
154    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
155        let mut txt = TXT::new();
156        for v in value.as_bytes().chunks(MAX_CHARACTER_STRING_LENGTH - 1) {
157            txt.add_char_string(CharacterString::new(v)?);
158        }
159        Ok(txt)
160    }
161}
162
163impl<'a> TryFrom<TXT<'a>> for String {
164    type Error = std::string::FromUtf8Error;
165
166    fn try_from(val: TXT<'a>) -> Result<Self, Self::Error> {
167        let init = Vec::with_capacity(val.len());
168
169        let bytes = val.strings.into_iter().fold(init, |mut acc, val| {
170            acc.extend(val.data.as_ref());
171            acc
172        });
173        String::from_utf8(bytes)
174    }
175}
176
177impl<'a> WireFormat<'a> for TXT<'a> {
178    const MINIMUM_LEN: usize = 1;
179
180    fn parse(data: &mut crate::bytes_buffer::BytesBuffer<'a>) -> crate::Result<Self>
181    where
182        Self: Sized,
183    {
184        let mut strings = Vec::new();
185        let mut size = 0;
186
187        while data.has_remaining() {
188            let char_str = CharacterString::parse(data)?;
189            size += char_str.len();
190            strings.push(char_str);
191        }
192
193        Ok(Self { strings, size })
194    }
195
196    fn len(&self) -> usize {
197        if self.strings.is_empty() {
198            Self::MINIMUM_LEN
199        } else {
200            self.size
201        }
202    }
203
204    fn write_to<T: std::io::Write>(&self, out: &mut T) -> crate::Result<()> {
205        if self.strings.is_empty() {
206            out.write_all(&[0])?;
207        } else {
208            for string in &self.strings {
209                string.write_to(out)?;
210            }
211        }
212        Ok(())
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use crate::{rdata::RData, ResourceRecord};
219    use std::convert::TryInto;
220
221    use super::*;
222
223    #[test]
224    pub fn parse_and_write_txt() -> Result<(), Box<dyn std::error::Error>> {
225        let mut out = vec![];
226        let txt = TXT::new()
227            .with_char_string("version=0.1".try_into()?)
228            .with_char_string("proto=123".try_into()?);
229
230        txt.write_to(&mut out)?;
231        assert_eq!(out.len(), txt.len());
232
233        let txt2 = TXT::parse(&mut out[..].into())?;
234        assert_eq!(2, txt2.strings.len());
235        assert_eq!(txt.strings[0], txt2.strings[0]);
236        assert_eq!(txt.strings[1], txt2.strings[1]);
237
238        Ok(())
239    }
240
241    #[test]
242    pub fn get_attributes() -> Result<(), Box<dyn std::error::Error>> {
243        let attributes = TXT::new()
244            .with_string("version=0.1")?
245            .with_string("flag")?
246            .with_string("with_eq=eq=")?
247            .with_string("version=dup")?
248            .with_string("empty=")?
249            .attributes();
250
251        assert_eq!(4, attributes.len());
252        assert_eq!(Some("0.1".to_owned()), attributes["version"]);
253        assert_eq!(Some("eq=".to_owned()), attributes["with_eq"]);
254        assert_eq!(Some(String::new()), attributes["empty"]);
255        assert_eq!(None, attributes["flag"]);
256
257        Ok(())
258    }
259
260    #[test]
261    fn parse_sample() -> Result<(), Box<dyn std::error::Error>> {
262        let sample_file = std::fs::read("samples/zonefile/TXT.sample")?;
263
264        let sample_rdata = match ResourceRecord::parse(&mut sample_file[..].into())?.rdata {
265            RData::TXT(rdata) => rdata,
266            _ => unreachable!(),
267        };
268
269        let strings = vec!["\"foo\nbar\"".try_into()?];
270        assert_eq!(sample_rdata.strings, strings);
271
272        Ok(())
273    }
274
275    #[test]
276    fn write_and_parse_large_txt() -> Result<(), Box<dyn std::error::Error>> {
277        let string = "X".repeat(1000);
278        let txt: TXT = string.as_str().try_into()?;
279
280        let mut bytes = Vec::new();
281        assert!(txt.write_to(&mut bytes).is_ok());
282
283        let parsed_txt = TXT::parse(&mut bytes[..].into())?;
284        let parsed_string: String = parsed_txt.try_into()?;
285
286        assert_eq!(parsed_string, string);
287
288        Ok(())
289    }
290
291    #[test]
292    fn write_and_parse_large_attributes() -> Result<(), Box<dyn std::error::Error>> {
293        let big_value = "f".repeat(1000);
294
295        let string = format!("foo={};;flag;bar={}", big_value, big_value);
296        let txt: TXT = string.as_str().try_into()?;
297        let attributes = txt.long_attributes()?;
298
299        assert_eq!(Some(big_value.to_owned()), attributes["bar"]);
300
301        Ok(())
302    }
303}
simple_dns/dns/rdata/txt.rs

simple_dns/dns/rdata/
txt.rs
