use std::{
collections::HashMap,
convert::{TryFrom, TryInto},
};
use crate::dns::{PacketPart, MAX_CHARACTER_STRING_LENGTH};
use crate::CharacterString;
use super::RR;
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub struct TXT<'a> {
strings: Vec<CharacterString<'a>>,
size: usize,
}
impl<'a> RR for TXT<'a> {
const TYPE_CODE: u16 = 16;
}
impl<'a> Default for TXT<'a> {
fn default() -> Self {
Self::new()
}
}
impl<'a> TXT<'a> {
pub fn new() -> Self {
Self {
strings: vec![],
size: 0,
}
}
pub fn add_string(&mut self, char_string: &'a str) -> crate::Result<()> {
self.add_char_string(char_string.try_into()?);
Ok(())
}
pub fn add_char_string(&mut self, char_string: CharacterString<'a>) {
self.size += char_string.len();
self.strings.push(char_string);
}
pub fn with_string(mut self, char_string: &'a str) -> crate::Result<Self> {
self.add_char_string(char_string.try_into()?);
Ok(self)
}
pub fn with_char_string(mut self, char_string: CharacterString<'a>) -> Self {
self.add_char_string(char_string);
self
}
pub fn attributes(&self) -> HashMap<String, Option<String>> {
let mut attributes = HashMap::new();
for char_str in &self.strings {
let mut splited = char_str.data.splitn(2, |c| *c == b'=');
let key = match splited.next() {
Some(key) => match std::str::from_utf8(key) {
Ok(key) => key.to_owned(),
Err(_) => continue,
},
None => continue,
};
let value = match splited.next() {
Some(value) if !value.is_empty() => match std::str::from_utf8(value) {
Ok(v) => Some(v.to_owned()),
Err(_) => Some(String::new()),
},
Some(_) => Some(String::new()),
_ => None,
};
attributes.entry(key).or_insert(value);
}
attributes
}
pub fn long_attributes(self) -> crate::Result<HashMap<String, Option<String>>> {
let mut attributes = HashMap::new();
let full_string: String = match self.try_into() {
Ok(string) => string,
Err(err) => return Err(crate::SimpleDnsError::InvalidUtf8String(err)),
};
let parts = full_string.split(|c| (c as u8) == b';');
for part in parts {
let key_value = part.splitn(2, |c| (c as u8) == b'=').collect::<Vec<&str>>();
let key = key_value[0];
let value = match key_value.len() > 1 {
true => Some(key_value[1].to_owned()),
_ => None,
};
if !key.is_empty() {
attributes.entry(key.to_owned()).or_insert(value);
}
}
Ok(attributes)
}
pub fn into_owned<'b>(self) -> TXT<'b> {
TXT {
strings: self.strings.into_iter().map(|s| s.into_owned()).collect(),
size: self.size,
}
}
}
impl<'a> TryFrom<HashMap<String, Option<String>>> for TXT<'a> {
type Error = crate::SimpleDnsError;
fn try_from(value: HashMap<String, Option<String>>) -> Result<Self, Self::Error> {
let mut txt = TXT::new();
for (key, value) in value {
match value {
Some(value) => {
txt.add_char_string(format!("{}={}", &key, &value).try_into()?);
}
None => txt.add_char_string(key.try_into()?),
}
}
Ok(txt)
}
}
impl<'a> TryFrom<&'a str> for TXT<'a> {
type Error = crate::SimpleDnsError;
fn try_from(value: &'a str) -> Result<Self, Self::Error> {
let mut txt = TXT::new();
for v in value.as_bytes().chunks(MAX_CHARACTER_STRING_LENGTH - 1) {
txt.add_char_string(CharacterString::new(v)?);
}
Ok(txt)
}
}
impl<'a> TryFrom<TXT<'a>> for String {
type Error = std::string::FromUtf8Error;
fn try_from(val: TXT<'a>) -> Result<Self, Self::Error> {
let init = Vec::with_capacity(val.len());
let bytes = val.strings.into_iter().fold(init, |mut acc, val| {
acc.extend(val.data.as_ref());
acc
});
String::from_utf8(bytes)
}
}
impl<'a> PacketPart<'a> for TXT<'a> {
fn parse(data: &'a [u8], position: usize) -> crate::Result<Self>
where
Self: Sized,
{
let mut strings = Vec::new();
let mut curr_position = position;
while curr_position < data.len() {
let char_str = CharacterString::parse(&data[curr_position..], 0)?;
curr_position += char_str.len();
strings.push(char_str);
}
Ok(Self {
strings,
size: data.len() - position,
})
}
fn len(&self) -> usize {
if self.strings.is_empty() {
1
} else {
self.size
}
}
fn write_to<T: std::io::Write>(&self, out: &mut T) -> crate::Result<()> {
if self.strings.is_empty() {
out.write_all(&[0])?;
} else {
for string in &self.strings {
string.write_to(out)?;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::{rdata::RData, ResourceRecord};
use std::convert::TryInto;
use super::*;
#[test]
pub fn parse_and_write_txt() -> Result<(), Box<dyn std::error::Error>> {
let mut out = vec![];
let txt = TXT::new()
.with_char_string("version=0.1".try_into()?)
.with_char_string("proto=123".try_into()?);
txt.write_to(&mut out)?;
assert_eq!(out.len(), txt.len());
let txt2 = TXT::parse(&out, 0)?;
assert_eq!(2, txt2.strings.len());
assert_eq!(txt.strings[0], txt2.strings[0]);
assert_eq!(txt.strings[1], txt2.strings[1]);
Ok(())
}
#[test]
pub fn get_attributes() -> Result<(), Box<dyn std::error::Error>> {
let attributes = TXT::new()
.with_string("version=0.1")?
.with_string("flag")?
.with_string("with_eq=eq=")?
.with_string("version=dup")?
.with_string("empty=")?
.attributes();
assert_eq!(4, attributes.len());
assert_eq!(Some("0.1".to_owned()), attributes["version"]);
assert_eq!(Some("eq=".to_owned()), attributes["with_eq"]);
assert_eq!(Some(String::new()), attributes["empty"]);
assert_eq!(None, attributes["flag"]);
Ok(())
}
#[test]
fn parse_sample() -> Result<(), Box<dyn std::error::Error>> {
let sample_file = std::fs::read("samples/zonefile/TXT.sample")?;
let sample_rdata = match ResourceRecord::parse(&sample_file, 0)?.rdata {
RData::TXT(rdata) => rdata,
_ => unreachable!(),
};
let strings = vec!["\"foo\nbar\"".try_into()?];
assert_eq!(sample_rdata.strings, strings);
Ok(())
}
#[test]
fn write_and_parse_large_txt() -> Result<(), Box<dyn std::error::Error>> {
let string = "X".repeat(1000);
let txt: TXT = string.as_str().try_into()?;
let mut bytes = Vec::new();
assert!(txt.write_to(&mut bytes).is_ok());
let parsed_txt = TXT::parse(&bytes, 0)?;
let parsed_string: String = parsed_txt.try_into()?;
assert_eq!(parsed_string, string);
Ok(())
}
#[test]
fn write_and_parse_large_attributes() -> Result<(), Box<dyn std::error::Error>> {
let big_value = "f".repeat(1000);
let string = format!("foo={};;flag;bar={}", big_value, big_value);
let txt: TXT = string.as_str().try_into()?;
let attributes = txt.long_attributes()?;
assert_eq!(Some(big_value.to_owned()), attributes["bar"]);
Ok(())
}
}