webrtc_mdns/message/
builder.rs

1use std::collections::HashMap;
2
3use super::header::*;
4use super::question::*;
5use super::resource::*;
6use super::*;
7use crate::error::*;
8
9// A Builder allows incrementally packing a DNS message.
10//
11// Example usage:
12//	b := NewBuilder(Header{...})
13//	b.enable_compression()
14//	// Optionally start a section and add things to that section.
15//	// Repeat adding sections as necessary.
16//	buf, err := b.Finish()
17//	// If err is nil, buf[2:] will contain the built bytes.
18#[derive(Default)]
19pub struct Builder {
20    // msg is the storage for the message being built.
21    pub msg: Option<Vec<u8>>,
22
23    // section keeps track of the current section being built.
24    pub section: Section,
25
26    // header keeps track of what should go in the header when Finish is
27    // called.
28    pub header: HeaderInternal,
29
30    // start is the starting index of the bytes allocated in msg for header.
31    pub start: usize,
32
33    // compression is a mapping from name suffixes to their starting index
34    // in msg.
35    pub compression: Option<HashMap<String, usize>>,
36}
37
38impl Builder {
39    // NewBuilder creates a new builder with compression disabled.
40    //
41    // Note: Most users will want to immediately enable compression with the
42    // enable_compression method. See that method's comment for why you may or may
43    // not want to enable compression.
44    //
45    // The DNS message is appended to the provided initial buffer buf (which may be
46    // nil) as it is built. The final message is returned by the (*Builder).Finish
47    // method, which may return the same underlying array if there was sufficient
48    // capacity in the slice.
49    pub fn new(h: &Header) -> Self {
50        let (id, bits) = h.pack();
51
52        Builder {
53            msg: Some(vec![0; HEADER_LEN]),
54            start: 0,
55            section: Section::Header,
56            header: HeaderInternal {
57                id,
58                bits,
59                ..Default::default()
60            },
61            compression: None,
62        }
63
64        //var hb [HEADER_LEN]byte
65        //b.msg = append(b.msg, hb[:]...)
66        //return b
67    }
68
69    // enable_compression enables compression in the Builder.
70    //
71    // Leaving compression disabled avoids compression related allocations, but can
72    // result in larger message sizes. Be careful with this mode as it can cause
73    // messages to exceed the UDP size limit.
74    //
75    // According to RFC 1035, section 4.1.4, the use of compression is optional, but
76    // all implementations must accept both compressed and uncompressed DNS
77    // messages.
78    //
79    // Compression should be enabled before any sections are added for best results.
80    pub fn enable_compression(&mut self) {
81        self.compression = Some(HashMap::new());
82    }
83
84    fn start_check(&self, section: Section) -> Result<()> {
85        if self.section <= Section::NotStarted {
86            return Err(Error::ErrNotStarted);
87        }
88        if self.section > section {
89            return Err(Error::ErrSectionDone);
90        }
91
92        Ok(())
93    }
94
95    // start_questions prepares the builder for packing Questions.
96    pub fn start_questions(&mut self) -> Result<()> {
97        self.start_check(Section::Questions)?;
98        self.section = Section::Questions;
99        Ok(())
100    }
101
102    // start_answers prepares the builder for packing Answers.
103    pub fn start_answers(&mut self) -> Result<()> {
104        self.start_check(Section::Answers)?;
105        self.section = Section::Answers;
106        Ok(())
107    }
108
109    // start_authorities prepares the builder for packing Authorities.
110    pub fn start_authorities(&mut self) -> Result<()> {
111        self.start_check(Section::Authorities)?;
112        self.section = Section::Authorities;
113        Ok(())
114    }
115
116    // start_additionals prepares the builder for packing Additionals.
117    pub fn start_additionals(&mut self) -> Result<()> {
118        self.start_check(Section::Additionals)?;
119        self.section = Section::Additionals;
120        Ok(())
121    }
122
123    fn increment_section_count(&mut self) -> Result<()> {
124        let section = self.section;
125        let (count, err) = match section {
126            Section::Questions => (&mut self.header.questions, Error::ErrTooManyQuestions),
127            Section::Answers => (&mut self.header.answers, Error::ErrTooManyAnswers),
128            Section::Authorities => (&mut self.header.authorities, Error::ErrTooManyAuthorities),
129            Section::Additionals => (&mut self.header.additionals, Error::ErrTooManyAdditionals),
130            Section::NotStarted => return Err(Error::ErrNotStarted),
131            Section::Done => return Err(Error::ErrSectionDone),
132            Section::Header => return Err(Error::ErrSectionHeader),
133        };
134
135        if *count == u16::MAX {
136            Err(err)
137        } else {
138            *count += 1;
139            Ok(())
140        }
141    }
142
143    // question adds a single question.
144    pub fn add_question(&mut self, q: &Question) -> Result<()> {
145        if self.section < Section::Questions {
146            return Err(Error::ErrNotStarted);
147        }
148        if self.section > Section::Questions {
149            return Err(Error::ErrSectionDone);
150        }
151        let msg = self.msg.take();
152        if let Some(mut msg) = msg {
153            msg = q.pack(msg, &mut self.compression, self.start)?;
154            self.increment_section_count()?;
155            self.msg = Some(msg);
156        }
157
158        Ok(())
159    }
160
161    fn check_resource_section(&self) -> Result<()> {
162        if self.section < Section::Answers {
163            return Err(Error::ErrNotStarted);
164        }
165        if self.section > Section::Additionals {
166            return Err(Error::ErrSectionDone);
167        }
168        Ok(())
169    }
170
171    // Resource adds a single resource.
172    pub fn add_resource(&mut self, r: &mut Resource) -> Result<()> {
173        self.check_resource_section()?;
174
175        if let Some(body) = &r.body {
176            r.header.typ = body.real_type();
177        } else {
178            return Err(Error::ErrNilResourceBody);
179        }
180
181        if let Some(msg) = self.msg.take() {
182            let (mut msg, len_off) = r.header.pack(msg, &mut self.compression, self.start)?;
183            let pre_len = msg.len();
184            if let Some(body) = &r.body {
185                msg = body.pack(msg, &mut self.compression, self.start)?;
186                r.header.fix_len(&mut msg, len_off, pre_len)?;
187                self.increment_section_count()?;
188            }
189            self.msg = Some(msg);
190        }
191
192        Ok(())
193    }
194
195    // Finish ends message building and generates a binary message.
196    pub fn finish(&mut self) -> Result<Vec<u8>> {
197        if self.section < Section::Header {
198            return Err(Error::ErrNotStarted);
199        }
200        self.section = Section::Done;
201
202        // Space for the header was allocated in NewBuilder.
203        let buf = self.header.pack(vec![]);
204        assert_eq!(buf.len(), HEADER_LEN);
205        if let Some(mut msg) = self.msg.take() {
206            msg[..HEADER_LEN].copy_from_slice(&buf[..HEADER_LEN]);
207            Ok(msg)
208        } else {
209            Err(Error::ErrEmptyBuilderMsg)
210        }
211    }
212}
webrtc_mdns/message/builder.rs

webrtc_mdns/message/
builder.rs
