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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};
use {Opcode, ResponseCode, Header, QueryType, QueryClass};
/// Allows to build a DNS packet
///
/// Both query and answer packets may be built with this interface, although,
/// much of functionality is not implemented yet.
#[derive(Debug)]
pub struct Builder {
buf: Vec<u8>,
}
impl Builder {
/// Creates a new query
///
/// Initially all sections are empty. You're expected to fill
/// the questions section with `add_question`
pub fn new_query(id: u16, recursion: bool) -> Builder {
let mut buf = Vec::with_capacity(512);
let head = Header {
id: id,
query: true,
opcode: Opcode::StandardQuery,
authoritative: false,
truncated: false,
recursion_desired: recursion,
recursion_available: false,
authenticated_data: false,
checking_disabled: false,
response_code: ResponseCode::NoError,
questions: 0,
answers: 0,
nameservers: 0,
additional: 0,
};
buf.extend([0u8; 12].iter());
head.write(&mut buf[..12]);
Builder { buf: buf }
}
/// Adds a question to the packet
///
/// # Panics
///
/// * Answers, nameservers or additional section has already been written
/// * There are already 65535 questions in the buffer.
/// * When name is invalid
pub fn add_question(&mut self, qname: &str,
qtype: QueryType, qclass: QueryClass)
-> &mut Builder
{
if &self.buf[6..12] != b"\x00\x00\x00\x00\x00\x00" {
panic!("Too late to add a question");
}
self.write_name(qname);
self.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
self.buf.write_u16::<BigEndian>(qclass as u16).unwrap();
let oldq = BigEndian::read_u16(&self.buf[4..6]);
if oldq == 65535 {
panic!("Too many questions");
}
BigEndian::write_u16(&mut self.buf[4..6], oldq+1);
self
}
fn write_name(&mut self, name: &str) {
for part in name.split('.') {
assert!(part.len() < 63);
let ln = part.len() as u8;
self.buf.push(ln);
self.buf.extend(part.as_bytes());
}
self.buf.push(0);
}
/// Returns the final packet
///
/// When packet is not truncated method returns `Ok(packet)`. If
/// packet is truncated the method returns `Err(packet)`. In both
/// cases the packet is fully valid.
///
/// In the server implementation you may use
/// `x.build().unwrap_or_else(|x| x)`.
///
/// In the client implementation it's probably unwise to send truncated
/// packet, as it doesn't make sense. Even panicking may be more
/// appropriate.
// TODO(tailhook) does the truncation make sense for TCP, and how
// to treat it for EDNS0?
pub fn build(mut self) -> Result<Vec<u8>,Vec<u8>> {
// TODO(tailhook) optimize labels
if self.buf.len() > 512 {
Header::set_truncated(&mut self.buf[..12]);
Err(self.buf)
} else {
Ok(self.buf)
}
}
}
#[cfg(test)]
mod test {
use QueryType as QT;
use QueryClass as QC;
use super::Builder;
#[test]
fn build_query() {
let mut bld = Builder::new_query(1573, true);
bld.add_question("example.com", QT::A, QC::IN);
let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
\x07example\x03com\x00\x00\x01\x00\x01";
assert_eq!(&bld.build().unwrap()[..], &result[..]);
}
#[test]
fn build_srv_query() {
let mut bld = Builder::new_query(23513, true);
bld.add_question("_xmpp-server._tcp.gmail.com", QT::SRV, QC::IN);
let result = b"[\xd9\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
\x0c_xmpp-server\x04_tcp\x05gmail\x03com\x00\x00!\x00\x01";
assert_eq!(&bld.build().unwrap()[..], &result[..]);
}
}