1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

use bytes::{Bytes, BytesMut};

use crate::{Dns, Flags};

use num_traits::ToPrimitive;

use super::{encode_u16, encode_u8, EncodeError, EncodeResult};

use std::collections::HashMap;

impl Flags {
    pub fn encode(&self, bytes: &mut BytesMut) -> EncodeResult {
        let mut buffer = 0u8;
        if self.qr {
            buffer |= 0b1000_0000;
        }
        if let Some(opcode) = self.opcode.to_u8() {
            buffer |= opcode << 3;
        } else {
            return Err(EncodeError::OpcodeError);
        }
        if self.aa {
            buffer |= 0b0000_0100;
        }
        if self.tc {
            buffer |= 0b0000_0010;
        }
        if self.rd {
            buffer |= 0b0000_0001;
        }
        encode_u8(bytes, buffer);

        buffer = 0;
        if self.ra {
            buffer |= 0b1000_0000;
        }
        if self.ad {
            buffer |= 0b0010_0000;
        }
        if self.cd {
            buffer |= 0b0001_0000;
        }
        if let Some(rcode) = self.rcode.to_u8() {
            buffer |= rcode;
        } else {
            return Err(EncodeError::RCodeError);
        }
        encode_u8(bytes, buffer);

        Ok(())
    }
}

impl Dns {
    pub fn encode(&self, bytes: &mut BytesMut) -> EncodeResult {
        encode_u16(bytes, self.id);
        self.flags.encode(bytes)?;
        encode_u16(bytes, self.questions.len() as u16);
        encode_u16(bytes, self.answers.len() as u16);
        encode_u16(bytes, self.authorities.len() as u16);
        encode_u16(bytes, self.additionals.len() as u16);

        let mut compression = HashMap::new();

        for question in &self.questions {
            question.encode(bytes, &mut compression)?;
        }

        for answer in &self.answers {
            answer.encode(bytes, &mut compression)?;
        }

        for authority in &self.authorities {
            authority.encode(bytes, &mut compression)?;
        }

        for additional in &self.additionals {
            additional.encode(bytes, &mut compression)?;
        }

        Ok(())
    }

    pub fn to_bytes(&self) -> Result<Bytes, EncodeError> {
        let mut bytes = BytesMut::new();
        self.encode(&mut bytes)?;
        Ok(bytes.freeze())
    }
}