Write answers

src/main.rs

@@ -2,8 +2,8 @@ mod message;
 mod parser;
 mod writer;
 use anyhow::Result;
-use message::{Message, Question};
-use std::net::UdpSocket;
+use message::{Answer, Answers, Data, Message, Question, Questions};
+use std::net::{Ipv4Addr, UdpSocket};
 
 fn main() -> Result<()> {
     println!("Logs from your program will appear here!");
@@ -27,7 +27,16 @@ fn main() -> Result<()> {
 
 fn look_up(query: Message) -> Result<Message> {
     let mut msg = Message::new_response(query);
+
     let q = Question::new_aa("codecrafters.io");
-    msg.set_questions(vec![q]);
+    let qs = Questions::try_from(vec![q])?;
+    msg.set_questions(qs);
+
+    let ip = Ipv4Addr::new(8, 8, 8, 8);
+    let data = Data::Ipv4(ip);
+    let a = Answer::new_aa("codecrafters.io".into(), data);
+    let ans = Answers::try_from(vec![a])?;
+    msg.set_answers(ans);
+
     Ok(msg)
 }

src/message.rs

@@ -1,7 +1,8 @@
 pub mod header;
+use anyhow::Result;
 pub use header::Header;
-use header::QueryMode;
-use std::ops::Deref;
+use header::{OpCode, QueryMode};
+use std::{net::Ipv4Addr, ops::Deref};
 
 #[repr(u16)]
 #[derive(Clone, Copy, Debug)]
@@ -41,29 +42,138 @@ impl Deref for Questions {
     }
 }
 
+impl TryFrom<Vec<Question>> for Questions {
+    type Error = anyhow::Error;
+    fn try_from(qs: Vec<Question>) -> Result<Self> {
+        anyhow::ensure!(
+            qs.len() <= std::u16::MAX as usize,
+            "Exceed supported number questions: {}",
+            qs.len()
+        );
+        Ok(Questions(qs))
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum Data {
+    Ipv4(Ipv4Addr),
+}
+
+impl Data {
+    fn len(&self) -> u16 {
+        match self {
+            Self::Ipv4(ip) => ip.octets().len() as u16,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Answer {
+    name: String,
+    record: Record,
+    class: Class,
+    ttl: u32,
+    length: u16,
+    data: Data,
+}
+
+impl Answer {
+    pub fn new_aa(name: String, data: Data) -> Self {
+        Self {
+            name,
+            record: Record::AA,
+            class: Class::IN,
+            ttl: 60,
+            length: data.len(),
+            data,
+        }
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn record(&self) -> Record {
+        self.record
+    }
+
+    pub fn class(&self) -> Class {
+        self.class
+    }
+
+    pub fn ttl(&self) -> u32 {
+        self.ttl
+    }
+
+    pub fn len(&self) -> u16 {
+        self.length
+    }
+
+    pub fn data(&self) -> &Data {
+        &self.data
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct Answers(Vec<Answer>);
+
+impl Deref for Answers {
+    type Target = Vec<Answer>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl TryFrom<Vec<Answer>> for Answers {
+    type Error = anyhow::Error;
+    fn try_from(qs: Vec<Answer>) -> Result<Self> {
+        anyhow::ensure!(
+            qs.len() <= std::u16::MAX as usize,
+            "Exceed supported number of Answers: {}",
+            qs.len()
+        );
+        Ok(Answers(qs))
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Message {
     header: Header,
     questions: Questions,
+    answers: Answers,
 }
 
 impl Message {
-    pub fn new(header: Header, questions: Vec<Question>) -> Self {
+    pub fn new_query(header: Header, questions: Vec<Question>) -> Self {
         Self {
             header,
             questions: Questions(questions),
+            answers: Default::default(),
         }
     }
 
     pub fn new_response(query: Message) -> Self {
+        let mut header = Header::response(query.header.id);
+        header.op_code = query.header.op_code;
+        header.rd = query.header.rd;
+        header.r_code = if query.header.op_code == OpCode(0) {
+            OpCode::no_error()
+        } else {
+            OpCode::not_implemented()
+        };
         Self {
-            header: Header::response(query.header.id),
-            questions: Questions::default(),
+            header,
+            questions: Default::default(),
+            answers: Default::default(),
         }
     }
 }
 
 impl Message {
+    pub fn header(&self) -> &Header {
+        &self.header
+    }
+
     pub fn is_query(&self) -> bool {
         self.header.qr == QueryMode::Query
     }
@@ -72,15 +182,17 @@ impl Message {
         &self.questions
     }
 
-    pub fn set_questions(&mut self, qs: Vec<Question>) {
-        if qs.len() > std::u16::MAX as usize {
-            panic!("Number of items exceeds supported len: {}", qs.len());
-        }
+    pub fn set_questions(&mut self, qs: Questions) {
         self.header.qd_count = qs.len() as u16;
-        self.questions = Questions(qs);
+        self.questions = qs;
     }
 
-    pub fn header(&self) -> &Header {
-        &self.header
+    pub fn answers(&self) -> &Answers {
+        &self.answers
+    }
+
+    pub fn set_answers(&mut self, ans: Answers) {
+        self.header.an_count = ans.len() as u16;
+        self.answers = ans;
     }
 }

src/message/header.rs

@@ -21,6 +21,16 @@ impl From<u8> for QueryMode {
 #[derive(Clone, Copy, Debug, PartialEq, Default)]
 pub struct OpCode(pub u8); // 4 bits
 
+impl OpCode {
+    pub fn no_error() -> Self {
+        OpCode(0)
+    }
+
+    pub fn not_implemented() -> Self {
+        OpCode(4)
+    }
+}
+
 #[repr(u8)]
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub enum Authoritative {
@@ -82,8 +92,8 @@ pub struct Header {
     pub op_code: OpCode,
     pub aa: Authoritative,
     pub tc: Truncation,
-    pub ra: Recursion,
     pub rd: Recursion,
+    pub ra: Recursion,
     pub z: Reserved,
     pub r_code: OpCode,
     pub qd_count: u16,

src/parser.rs

@@ -13,7 +13,7 @@ fn take_packet_id(i: &[u8]) -> ByteResult<PacketId> {
 }
 
 fn take_opcode(i: BitInput) -> BitResult<OpCode> {
-    map(bits::complete::take(1u8), |bits: u8| OpCode(bits)).parse(i)
+    map(bits::complete::take(4u8), |bits: u8| OpCode(bits)).parse(i)
 }
 
 fn take_enum<T: From<u8>>(i: BitInput) -> BitResult<T> {
@@ -49,8 +49,8 @@ fn parse_header(i: &[u8]) -> ByteResult<Header> {
         op_code: flags.1,
         aa: flags.2,
         tc: flags.3,
-        ra: flags.4,
-        rd: flags.5,
+        rd: flags.4,
+        ra: flags.5,
         z: flags.6,
         r_code: flags.7,
         qd_count,
@@ -70,7 +70,7 @@ fn parse_questions(i: &[u8], c: u16) -> ByteResult<Vec<Question>> {
 fn parse_message(i: &[u8]) -> ByteResult<Message> {
     let (i, header) = parse_header(i)?;
     let (i, questions) = parse_questions(i, header.qd_count)?;
-    let msg = Message::new(header, questions);
+    let msg = Message::new_query(header, questions);
     Ok((i, msg))
 }
 

src/writer.rs

@@ -1,20 +1,41 @@
 use crate::{
     message::{
         header::{Authoritative, QueryMode, Recursion, Truncation},
-        Header, Question, Questions,
+        Answer, Data, Header, Question,
     },
     Message,
 };
 use bytes::{BufMut, BytesMut};
 
-impl Header {
-    pub fn write(&self, buf: &mut BytesMut) {
+trait Writer {
+    fn write(&self, buf: &mut BytesMut);
+}
+
+impl Writer for &str {
+    fn write(&self, buf: &mut BytesMut) {
+        self.split('.').for_each(|l| {
+            buf.put_u8(l.len() as u8);
+            buf.put(l.as_bytes())
+        });
+        buf.put_u8(0);
+    }
+}
+
+impl<T: Writer> Writer for Vec<T> {
+    fn write(&self, buf: &mut BytesMut) {
+        self.iter().for_each(|i| i.write(buf))
+    }
+}
+
+impl Writer for Header {
+    fn write(&self, buf: &mut BytesMut) {
         buf.put_u16(self.id.0);
 
         let mut flags = 0u8;
         if self.qr == QueryMode::Response {
             flags |= 0b1_0000000;
         }
+        flags |= 0b01111000 & (self.op_code.0 << 3);
         if self.aa == Authoritative::Owned {
             flags |= 0b0000_0100;
         }
@@ -40,31 +61,39 @@ impl Header {
     }
 }
 
-impl Question {
-    pub fn write(&self, buf: &mut BytesMut) {
-        self.name.split('.').for_each(|l| {
-            buf.put_u8(l.len() as u8);
-            buf.put(l.as_bytes())
-        });
-        buf.put_u8(0);
+impl Writer for Question {
+    fn write(&self, buf: &mut BytesMut) {
+        self.name.as_str().write(buf);
         buf.put_u16(self.record as u16);
         buf.put_u16(self.class as u16);
     }
 }
 
-impl Questions {
-    pub fn write(&self, buf: &mut BytesMut) {
-        for q in self.iter() {
-            q.write(buf)
+impl Writer for Data {
+    fn write(&self, buf: &mut BytesMut) {
+        match self {
+            Self::Ipv4(ip) => buf.put_slice(&ip.octets()),
         }
     }
 }
 
+impl Writer for Answer {
+    fn write(&self, buf: &mut BytesMut) {
+        self.name().write(buf);
+        buf.put_u16(self.record() as u16);
+        buf.put_u16(self.class() as u16);
+        buf.put_u32(self.ttl());
+        buf.put_u16(self.len());
+        self.data().write(buf);
+    }
+}
+
 impl Message {
     pub fn flush(&self) -> BytesMut {
         let mut buf = BytesMut::with_capacity(12);
         self.header().write(&mut buf);
         self.questions().write(&mut buf);
+        self.answers().write(&mut buf);
         buf
     }
 }