fix: widcard problem, message compression strikes again (#5)

zns-daemon/src/db/models.rs

@@ -3,7 +3,7 @@ use diesel::sql_types::Text;
 use zns::{
     errors::ZNSError,
     labelstring::LabelString,
-    structs::{Class, Type, RR},
+    structs::{Class, RData, Type, RR},
 };
 
 use self::schema::records::{self};
@@ -109,7 +109,7 @@ pub fn insert_into_database(rr: &RR, connection: &mut PgConnection) -> Result<()
         class: rr.class.clone().into(),
         ttl: rr.ttl,
         rdlength: rr.rdlength as i32,
-        rdata: rr.rdata.clone(),
+        rdata: rr.rdata.clone().into(),
     };
 
     Record::create(connection, record).map_err(|e| ZNSError::Servfail {
@@ -137,13 +137,17 @@ pub fn get_from_database(
 
     Ok(records
         .into_iter()
-        .map(|record| RR {
-            name: LabelString::from(&record.name),
-            _type: Type::from(record._type as u16),
-            class: Class::from(record.class as u16),
-            ttl: record.ttl,
-            rdlength: record.rdlength as u16,
-            rdata: record.rdata,
+        .filter_map(|record| {
+            RData::from_safe(&record.rdata, &Type::from(record._type as u16))
+                .map(|rdata| RR {
+                    name: LabelString::from(&record.name),
+                    _type: Type::from(record._type as u16),
+                    class: Class::from(record.class as u16),
+                    ttl: record.ttl,
+                    rdlength: record.rdlength as u16,
+                    rdata,
+                })
+                .ok()
         })
         .collect())
 }
@@ -202,7 +206,7 @@ mod tests {
             &rr.name,
             Some(rr._type.clone()),
             rr.class.clone(),
-            Some(rr.rdata.clone()),
+            Some(rr.rdata.clone().into()),
             &mut connection,
         );
 

zns-daemon/src/handlers/query.rs

@@ -4,7 +4,7 @@ use zns::{
     errors::ZNSError,
     labelstring::LabelString,
     parser::ToBytes,
-    structs::{Class, Message, Question, RRClass, RRType, SoaRData, Type, RR},
+    structs::{Class, Message, Question, RData, RRClass, RRType, SoaRData, Type, RR},
 };
 
 use crate::{config::Config, db::models::get_from_database};
@@ -76,8 +76,8 @@ impl ResponseHandler for QueryHandler {
 }
 
 fn try_wildcard(question: &Question, connection: &mut PgConnection) -> Result<Vec<RR>, ZNSError> {
-    let qname = question.qname.clone().to_vec();
-    qname.to_vec()[0] = String::from("*");
+    let mut qname = question.qname.clone().to_vec();
+    qname[0] = String::from("*");
     Ok(get_from_database(
         &qname.into(),
         Some(question.qtype.clone()),
@@ -87,6 +87,7 @@ fn try_wildcard(question: &Question, connection: &mut PgConnection) -> Result<Ve
     .into_iter()
     .map(|mut rr| {
         rr.name.clone_from(&question.qname);
+        println!("{:#?}", rr);
         rr
     })
     .collect())
@@ -129,7 +130,7 @@ fn get_soa(name: &LabelString) -> Result<RR, ZNSError> {
         class: Class::Class(RRClass::IN),
         ttl: 11200,
         rdlength: 0,
-        rdata: SoaRData::to_bytes(rdata),
+        rdata: RData::from_safe(&SoaRData::to_bytes(rdata), &Type::Type(RRType::SOA))?,
     })
 }
 
@@ -165,4 +166,35 @@ mod tests {
         assert_eq!(result.answer[0], rr);
         assert_eq!(result.answer[1], rr);
     }
+
+    #[tokio::test]
+    async fn test_wildcard_query() {
+        let mut connection = get_test_connection();
+
+        let wildcard = Config::get().authoritative_zone.prepend("*".to_string());
+        let non_existent = Config::get()
+            .authoritative_zone
+            .prepend("nonexistent".to_string());
+
+        let mut rr = get_rr(Some(wildcard));
+
+        let mut message = get_message(Some(non_existent.clone()));
+        message.header.ancount = 0;
+        message.answer = vec![];
+
+        assert!(insert_into_database(&rr, &mut connection).is_ok());
+
+        rr.name = non_existent;
+
+        let result = QueryHandler::handle(
+            &message,
+            &Message::to_bytes(message.clone()),
+            &mut connection,
+        )
+        .await
+        .unwrap();
+        assert_eq!(result.header.ancount, 2);
+        assert_eq!(result.answer.len(), 2);
+        assert_eq!(result.answer[0], rr);
+    }
 }

zns-daemon/src/handlers/update/authenticate.rs

@@ -81,7 +81,8 @@ async fn validate_dnskey(
     )?
     .iter()
     .any(|rr| {
-        let mut reader = Reader::new(&rr.rdata);
+        let data: Vec<u8> = rr.rdata.clone().into();
+        let mut reader = Reader::new(&data);
         DNSKeyRData::from_bytes(&mut reader).is_ok_and(|dnskey| match sig.verify_dnskey(dnskey) {
             Ok(value) => value,
             Err(e) => {

zns-daemon/src/handlers/update/mod.rs

@@ -124,7 +124,7 @@ impl ResponseHandler for UpdateHandler {
                     &rr.name,
                     Some(rr._type.clone()),
                     Class::Class(RRClass::IN),
-                    Some(rr.rdata.clone()),
+                    Some(rr.rdata.clone().into()),
                     connection,
                 )
             }

zns-daemon/src/handlers/update/sig.rs

@@ -77,7 +77,8 @@ impl Sig {
         let mut request = datagram[0..datagram.len() - 11 - rr.rdlength as usize].to_vec();
         request[11] -= 1; // Decrease arcount
 
-        let mut reader = Reader::new(&rr.rdata);
+        let data: Vec<u8> = rr.rdata.clone().into();
+        let mut reader = Reader::new(&data);
         let key_rdata = SigRData::from_bytes(&mut reader)?;
 
         let now = SystemTime::now()
@@ -99,7 +100,7 @@ impl Sig {
             });
         }
 
-        let mut raw_data = rr.rdata[0..rr.rdata.len() - key_rdata.signature.len()].to_vec();
+        let mut raw_data = data[0..data.len() - key_rdata.signature.len()].to_vec();
         raw_data.extend(request);
 
         Ok(Sig {

zns/Cargo.toml

@@ -12,6 +12,7 @@ base64 = "0.22.0"
 int-enum = "1.1"
 thiserror = "1.0"
 arbitrary = { version = "^1.3.2", optional = true, features = ["derive"] }
+rand = {version = "0.8.5"}
 
 [dev-dependencies]
 zns = { path = ".", features = ["test-utils"] }

zns/src/labelstring.rs

@@ -39,6 +39,13 @@ impl LabelString {
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    #[cfg(feature = "test-utils")]
+    pub fn prepend(&self, element: String) -> Self {
+        let mut vec = self.0.clone();
+        vec.insert(0, element);
+        LabelString(vec)
+    }
 }
 
 impl PartialEq for LabelString {

zns/src/parser.rs

@@ -4,7 +4,9 @@ use crate::{
     errors::ZNSError,
     labelstring::LabelString,
     reader::Reader,
-    structs::{Class, Header, Message, Opcode, Question, RRClass, RRType, SoaRData, Type, RR},
+    structs::{
+        Class, Header, Message, Opcode, Question, RData, RRClass, RRType, SoaRData, Type, RR,
+    },
 };
 
 type Result<T> = std::result::Result<T, ZNSError>;
@@ -198,6 +200,47 @@ impl ToBytes for Question {
     }
 }
 
+impl From<RData> for Vec<u8> {
+    fn from(value: RData) -> Self {
+        match value {
+            RData::LabelString(labelstring) => LabelString::to_bytes(labelstring),
+            RData::Vec(vec) => vec,
+        }
+    }
+}
+
+impl RData {
+    pub fn from(reader: &mut Reader, rdlength: u16, rr_type: &Type) -> Result<Self> {
+        match rr_type {
+            Type::Type(RRType::CNAME) => Ok(Self::LabelString(LabelString::from_bytes(reader)?)),
+            _ => {
+                let data = reader.read(rdlength as usize)?;
+                Ok(Self::Vec(data.to_vec()))
+            }
+        }
+    }
+
+    pub fn from_safe(data: &[u8], rr_type: &Type) -> Result<Self> {
+        match rr_type {
+            Type::Type(RRType::CNAME) => Ok(Self::LabelString(LabelString::from_bytes(
+                &mut Reader::new(data),
+            )?)),
+            _ => Ok(Self::Vec(data.to_vec())),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        match self {
+            RData::LabelString(labelstring) => labelstring.len(),
+            RData::Vec(v) => v.len(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
 impl FromBytes for RR {
     fn from_bytes(reader: &mut Reader) -> Result<Self> {
         let name = LabelString::from_bytes(reader)?;
@@ -213,11 +256,11 @@ impl FromBytes for RR {
         } else {
             Ok(RR {
                 name,
+                rdata: RData::from(reader, rdlength, &_type)?,
                 _type,
                 class,
                 ttl,
                 rdlength,
-                rdata: reader.read(rdlength as usize)?,
             })
         }
     }
@@ -226,11 +269,12 @@ impl FromBytes for RR {
 impl ToBytes for RR {
     fn to_bytes(rr: Self) -> Vec<u8> {
         let mut result = LabelString::to_bytes(rr.name);
+        let rdata: Vec<u8> = rr.rdata.into();
         result.extend(u16::to_be_bytes(rr._type.into()));
         result.extend(u16::to_be_bytes(rr.class.into()));
         result.extend(i32::to_be_bytes(rr.ttl.to_owned()));
-        result.extend(u16::to_be_bytes(rr.rdata.len() as u16));
-        result.extend(rr.rdata);
+        result.extend(u16::to_be_bytes(rdata.len() as u16));
+        result.extend(rdata);
         result
     }
 }

zns/src/structs.rs

@@ -14,6 +14,7 @@ pub enum Type {
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 pub enum RRType {
     A = 1,
+    CNAME = 5,
     SOA = 6,
     SIG = 24,
     DNSKEY = 48,
@@ -59,6 +60,13 @@ pub enum Opcode {
     UPDATE = 5,
 }
 
+#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
+#[derive(Debug, Clone, PartialEq)]
+pub enum RData {
+    LabelString(LabelString),
+    Vec(Vec<u8>),
+}
+
 #[derive(Debug, Clone, PartialEq)]
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 pub struct Question {
@@ -96,7 +104,7 @@ pub struct RR {
     pub class: Class,
     pub ttl: i32,
     pub rdlength: u16,
-    pub rdata: Vec<u8>,
+    pub rdata: RData,
 }
 
 pub struct SoaRData {