Message parser and write_to done, need tests...

src/op/message.rs

@@ -1,5 +1,6 @@
 use super::header::Header;
 use super::query::Query;
+use super::super::rr::resource::Record;
 
 /*
  * RFC 1035        Domain Implementation and Specification    November 1987
@@ -39,7 +40,7 @@ use super::query::Query;
  * question.
  */
 pub struct Message {
-  header: Header, queries: Vec<Query>, /*answer: Answer, authority: domain::Name, additional: Additional*/
+  header: Header, queries: Vec<Query>, answers: Vec<Record>, name_servers: Vec<Record>, additionals: Vec<Record>
 }
 
 impl Message {
@@ -47,19 +48,54 @@ impl Message {
     let header = Header::parse(data);
 
     // get the questions
-    let count: usize = header.getQueryCount() as usize;
+    let count = header.getQueryCount() as usize;
     let mut queries = Vec::with_capacity(count);
     for _ in 0 .. count {
       queries.push(Query::parse(data));
     }
 
-    // get the answers
-    // let count: usize = header.getAnswerCount() as usize;
-    // let mut answers = Vec::with_capacity(count);
-    // for _ in 0 .. count {
-    //   answers.push(Answer)
-    // }
+    // get all counts before header moves
+    let answer_count = header.getAnswerCount() as usize;
+    let name_server_count = header.getNameServerCount() as usize;
+    let additional_count = header.getAdditionalCount() as usize;
 
-    Message { header: header, queries: queries}
+    Message {
+      header: header,
+      queries: queries,
+      answers: Self::parse_records(data, answer_count),
+      name_servers: Self::parse_records(data, name_server_count),
+      additionals: Self::parse_records(data, additional_count),
+    }
+  }
+
+  fn parse_records(data: &mut Vec<u8>, count: usize) -> Vec<Record> {
+    let mut records: Vec<Record> = Vec::with_capacity(count);
+    for _ in 0 .. count {
+       records.push(Record::parse(data))
+    }
+    records
+  }
+
+  pub fn write_to(&self, buf: &mut Vec<u8>) {
+    self.header.write_to(buf);
+
+    assert_eq!(self.header.getQueryCount() as usize, self.queries.len());
+    assert_eq!(self.header.getAnswerCount() as usize, self.answers.len());
+    assert_eq!(self.header.getNameServerCount() as usize, self.name_servers.len());
+    assert_eq!(self.header.getAdditionalCount() as usize, self.additionals.len());
+
+    for q in &self.queries {
+      q.write_to(buf);
+    }
+
+    Self::write_records(buf, &self.answers);
+    Self::write_records(buf, &self.name_servers);
+    Self::write_records(buf, &self.additionals);
+  }
+
+  fn write_records(buf: &mut Vec<u8>, records: &Vec<Record>) {
+    for r in records {
+      r.write_to(buf);
+    }
   }
 }

src/rr/resource.rs

@@ -1,6 +1,3 @@
-//use std::collections::HashMap;
-//use std::error::Error;
-use std::string::FromUtf8Error;
 use std::io::Read;
 use std::iter;
 use std::slice::{IterMut,Iter};
@@ -67,6 +64,7 @@ use super::util;
  *                 For example, the if the TYPE is A and the CLASS is IN,
  *                 the RDATA field is a 4 octet ARPA Internet address.
  */
+#[derive(PartialEq, Debug)]
 pub struct Record {
   name_labels: domain::Name,
   rr_type: RecordType,
@@ -76,10 +74,47 @@ pub struct Record {
 }
 
 impl Record {
+  /**
+   * Creates a not very useful empty record, use the setters to build a more useful object
+   */
+  pub fn new() -> Record {
+    Record {
+      name_labels: domain::Name::new(),
+      rr_type: RecordType::ANY,
+      dns_class: DNSClass::ANY,
+      ttl: 0,
+      rdata: RData::NULL { anything: vec![] }
+    }
+  }
+
+  pub fn add_name(&mut self, label: String) -> &mut Self {
+    self.name_labels.add_label(label);
+    self
+  }
+
+  pub fn rr_type(&mut self, rr_type: RecordType) -> &mut Self {
+    self.rr_type = rr_type;
+    self
+  }
+
+  pub fn dns_class(&mut self, dns_class: DNSClass) -> &mut Self {
+    self.dns_class = dns_class;
+    self
+  }
+
+  pub fn ttl(&mut self, ttl: i32) -> &mut Self {
+    self.ttl = ttl;
+    self
+  }
+
+  pub fn rdata(&mut self, rdata: RData) -> &mut Self {
+    self.rdata = rdata;
+    self
+  }
 
   /// parse a resource record line example:
   ///  WARNING: the record_bytes is 100% consumed and destroyed in this parsing process
-  pub fn parse(data: &mut Vec<u8>) -> Result<Record, FromUtf8Error> {
+  pub fn parse(data: &mut Vec<u8>) -> Record {
     // NAME            an owner name, i.e., the name of the node to which this
     //                 resource record pertains.
     let name_labels: domain::Name = domain::Name::parse(data);
@@ -109,7 +144,7 @@ impl Record {
     //                according to the TYPE and CLASS of the resource record.
     let rdata = RData::parse(data, record_type, rd_length);
 
-    Ok(Record{ name_labels: name_labels, rr_type: record_type, dns_class: class, ttl: ttl, rdata: rdata })
+    Record{ name_labels: name_labels, rr_type: record_type, dns_class: class, ttl: ttl, rdata: rdata }
   }
 
   pub fn write_to(&self, buf: &mut Vec<u8>) {
@@ -133,3 +168,36 @@ impl Record {
     }
   }
 }
+
+#[cfg(test)]
+mod tests {
+  use std::net::Ipv4Addr;
+
+  use super::*;
+
+  use super::super::record_data::RData;
+  use super::super::record_type::RecordType;
+  use super::super::dns_class::DNSClass;
+  use super::super::domain;
+  use super::super::util;
+
+
+  #[test]
+  fn test_write_and_parse() {
+    let mut record = Record::new();
+    record.add_name("www".to_string()).add_name("example".to_string()).add_name("com".to_string());
+    record.rr_type(RecordType::A);
+    record.dns_class(DNSClass::IN);
+    record.ttl(5);
+    record.rdata(RData::A { address: Ipv4Addr::new(192, 168, 0, 1)});
+
+    let mut buf: Vec<u8> = Vec::new();
+    record.write_to(&mut buf);
+
+    buf.reverse(); // reverse the stream...
+
+    let got = Record::parse(&mut buf);
+
+    assert_eq!(got, record);
+  }
+}