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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
/*
 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Message metadata

use std::convert::From;

use super::op_code::OpCode;
use super::response_code::ResponseCode;
use error::*;
use serialize::binary::*;

/// Metadata for the `Message` struct.
///
/// [RFC 1035, DOMAIN NAMES - IMPLEMENTATION AND SPECIFICATION, November 1987](https://tools.ietf.org/html/rfc1035)
///
/// ```text
/// 4.1.1. Header section format
///
/// The header contains the following fields
///
///                                    1  1  1  1  1  1
///      0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |                      ID                       |
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |QR|   Opcode  |AA|TC|RD|RA|ZZ|AD|CD|   RCODE   |  /// AD and CD from RFC4035
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |                    QDCOUNT / ZCOUNT           |
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |                    ANCOUNT / PRCOUNT          |
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |                    NSCOUNT / UPCOUNT          |
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///     |                    ARCOUNT / ADCOUNT          |
///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
///
/// where
///
/// Z               Reserved for future use.  Must be zero in all queries
///                 and responses.
///
/// ```
///
#[derive(Clone, Debug, PartialEq, PartialOrd)]
pub struct Header {
    id: u16,
    message_type: MessageType,
    op_code: OpCode,
    authoritative: bool,
    truncation: bool,
    recursion_desired: bool,
    recursion_available: bool,
    authentic_data: bool,
    checking_disabled: bool,
    response_code: u8, /* ideally u4 */
    query_count: u16,
    answer_count: u16,
    name_server_count: u16,
    additional_count: u16,
}

/// Message types are either Query (also Update) or Response
#[derive(Debug, PartialEq, PartialOrd, Copy, Clone)]
pub enum MessageType {
    /// Queries are Client requests, these are either Queries or Updates
    Query,
    /// Response message from the Server or upstream Resolver
    Response,
}

impl Default for Header {
    fn default() -> Self {
        Header {
            id: 0,
            message_type: MessageType::Query,
            op_code: OpCode::Query,
            authoritative: false,
            truncation: false,
            recursion_desired: false,
            recursion_available: false,
            authentic_data: false,
            checking_disabled: false,
            response_code: 0,
            query_count: 0,
            answer_count: 0,
            name_server_count: 0,
            additional_count: 0,
        }
    }
}

impl Header {
    // TODO: we should make id, message_type and op_code all required and non-editable
    /// A default Header, not very useful.
    pub fn new() -> Self {
        Default::default()
    }

    /// Length of the header, always 12 bytes
    #[inline(always)]
    pub fn len() -> usize {
        12 /* this is always 12 bytes */
    }

    /// Sets the id of the message, for queries this shoudl be random.
    pub fn set_id(&mut self, id: u16) -> &mut Self {
        self.id = id;
        self
    }

    /// Sets the message type, Queries and Updates both use Query.
    pub fn set_message_type(&mut self, message_type: MessageType) -> &mut Self {
        self.message_type = message_type;
        self
    }

    /// Set the operation code for the message
    pub fn set_op_code(&mut self, op_code: OpCode) -> &mut Self {
        self.op_code = op_code;
        self
    }

    /// From the server is specifies that it is an authoritative reqponse.
    pub fn set_authoritative(&mut self, authoritative: bool) -> &mut Self {
        self.authoritative = authoritative;
        self
    }

    /// Specifies that the records were too large for the payload.
    ///
    /// See EDNS or TCP for resolutions to truncation.
    pub fn set_truncated(&mut self, truncated: bool) -> &mut Self {
        self.truncation = truncated;
        self
    }

    /// Specify that the resolver should recursiviley request data from upstream DNS nodes
    pub fn set_recursion_desired(&mut self, recursion_desired: bool) -> &mut Self {
        self.recursion_desired = recursion_desired;
        self
    }

    /// Specifies that recursion is available from this or the remote resolver
    pub fn set_recursion_available(&mut self, recursion_available: bool) -> &mut Self {
        self.recursion_available = recursion_available;
        self
    }

    /// Specifies that the data is authnetic, i.e. the resolver believes all data to be valid through DNSSec
    pub fn set_authentic_data(&mut self, authentic_data: bool) -> &mut Self {
        self.authentic_data = authentic_data;
        self
    }

    /// Used during recursive resolution to specified if a resolver should or should not validate DNSSec signatures
    pub fn set_checking_disabled(&mut self, checking_disabled: bool) -> &mut Self {
        self.checking_disabled = checking_disabled;
        self
    }

    /// The low responsed code (original response codes before EDNS extensions)
    pub fn set_response_code(&mut self, response_code: ResponseCode) -> &mut Self {
        self.response_code = response_code.low();
        self
    }

    /// Number or query records in the message
    pub fn set_query_count(&mut self, query_count: u16) -> &mut Self {
        self.query_count = query_count;
        self
    }

    /// Number of answer records in the message
    pub fn set_answer_count(&mut self, answer_count: u16) -> &mut Self {
        self.answer_count = answer_count;
        self
    }

    /// Number of name server records in the message
    pub fn set_name_server_count(&mut self, name_server_count: u16) -> &mut Self {
        self.name_server_count = name_server_count;
        self
    }

    /// Number of additional records in the message
    pub fn set_additional_count(&mut self, additional_count: u16) -> &mut Self {
        self.additional_count = additional_count;
        self
    }

    /// ```text
    /// ID              A 16 bit identifier assigned by the program that
    ///                 generates any kind of query.  This identifier is copied
    ///                 the corresponding reply and can be used by the requester
    ///                 to match up replies to outstanding queries.
    /// ```
    pub fn id(&self) -> u16 {
        self.id
    }

    /// ```text
    /// QR              A one bit field that specifies whether this message is a
    ///                 query (0), or a response (1).
    /// ```
    pub fn message_type(&self) -> MessageType {
        self.message_type
    }

    /// ```text
    /// OPCODE          A four bit field that specifies kind of query in this
    ///                 message.  This value is set by the originator of a query
    ///                 and copied into the response.  The values are: <see super::op_code>
    /// ```
    pub fn op_code(&self) -> OpCode {
        self.op_code
    }

    /// ```text
    /// AA              Authoritative Answer - this bit is valid in responses,
    ///                 and specifies that the responding name server is an
    ///                 authority for the domain name in question section.
    ///
    ///                 Note that the contents of the answer section may have
    ///                 multiple owner names because of aliases.  The AA bit
    ///                 corresponds to the name which matches the query name, or
    ///                 the first owner name in the answer section.
    /// ```
    pub fn authoritative(&self) -> bool {
        self.authoritative
    }

    /// ```text
    /// TC              TrunCation - specifies that this message was truncated
    ///                 due to length greater than that permitted on the
    ///                 transmission channel.
    /// ```
    pub fn truncated(&self) -> bool {
        self.truncation
    }

    /// ```text
    /// RD              Recursion Desired - this bit may be set in a query and
    ///                 is copied into the response.  If RD is set, it directs
    ///                 the name server to pursue the query recursively.
    ///                 Recursive query support is optional.
    /// ```
    pub fn recursion_desired(&self) -> bool {
        self.recursion_desired
    }

    /// ```text
    /// RA              Recursion Available - this be is set or cleared in a
    ///                 response, and denotes whether recursive query support is
    ///                 available in the name server.
    /// ```
    pub fn recursion_available(&self) -> bool {
        self.recursion_available
    }

    /// [RFC 4035, DNSSEC Resource Records, March 2005](https://tools.ietf.org/html/rfc4035#section-3.1.6)
    ///
    /// ```text
    ///
    /// 3.1.6.  The AD and CD Bits in an Authoritative Response
    ///
    ///   The CD and AD bits are designed for use in communication between
    ///   security-aware resolvers and security-aware recursive name servers.
    ///   These bits are for the most part not relevant to query processing by
    ///   security-aware authoritative name servers.
    ///
    ///   A security-aware name server does not perform signature validation
    ///   for authoritative data during query processing, even when the CD bit
    ///   is clear.  A security-aware name server SHOULD clear the CD bit when
    ///   composing an authoritative response.
    ///
    ///   A security-aware name server MUST NOT set the AD bit in a response
    ///   unless the name server considers all RRsets in the Answer and
    ///   Authority sections of the response to be authentic.  A security-aware
    ///   name server's local policy MAY consider data from an authoritative
    ///   zone to be authentic without further validation.  However, the name
    ///   server MUST NOT do so unless the name server obtained the
    ///   authoritative zone via secure means (such as a secure zone transfer
    ///   mechanism) and MUST NOT do so unless this behavior has been
    ///   configured explicitly.
    ///
    ///   A security-aware name server that supports recursion MUST follow the
    ///   rules for the CD and AD bits given in Section 3.2 when generating a
    ///   response that involves data obtained via recursion.
    /// ```
    pub fn authentic_data(&self) -> bool {
        self.authentic_data
    }

    /// see `is_authentic_data()`
    pub fn checking_disabled(&self) -> bool {
        self.checking_disabled
    }

    /// ```text
    /// RCODE           Response code - this 4 bit field is set as part of
    ///                 responses.  The values have the following
    ///                 interpretation: <see super::response_code>
    /// ```
    pub fn response_code(&self) -> u8 {
        self.response_code
    }

    /// ```text
    /// QDCOUNT         an unsigned 16 bit integer specifying the number of
    ///                 entries in the question section.
    /// ```
    ///
    /// # Return value
    ///
    /// If this is a query, this will return the number of queries in the query section of the
    //   message, fo updates this represents the zone count (must be no more than 1).
    pub fn query_count(&self) -> u16 {
        self.query_count
    }

    /// ```text
    /// ANCOUNT         an unsigned 16 bit integer specifying the number of
    ///                 resource records in the answer section.
    /// ```
    ///
    /// # Return value
    ///
    /// For query responses this is the number of records in the answer section, should be 0 for
    ///  requests, for updates this is the count of prerequisite records.
    pub fn answer_count(&self) -> u16 {
        self.answer_count
    }

    /// for queries this is the nameservers which are authorities for the SOA of the Record
    /// for updates this is the update record count
    /// ```text
    /// NSCOUNT         an unsigned 16 bit integer specifying the number of name
    ///                 server resource records in the authority records
    ///                 section.
    /// ```
    ///
    /// # Return value
    ///
    /// For query responses this is the number of authorities, or nameservers, in the name server
    ///  section, for updates this is the number of update records being sent.
    pub fn name_server_count(&self) -> u16 {
        self.name_server_count
    }

    /// ```text
    /// ARCOUNT         an unsigned 16 bit integer specifying the number of
    ///                 resource records in the additional records section.
    /// ```
    ///
    /// # Return value
    ///
    /// This is the additional record section count, this section may include EDNS options.
    pub fn additional_count(&self) -> u16 {
        self.additional_count
    }
}

impl BinEncodable for Header {
    fn emit(&self, encoder: &mut BinEncoder) -> ProtoResult<()> {
        encoder.reserve(12)?; // the 12 bytes for the following fields;

        // Id
        encoder.emit_u16(self.id)?;

        // IsQuery, OpCode, Authoritative, Truncation, RecursionDesired
        let mut q_opcd_a_t_r: u8 = if let MessageType::Response = self.message_type {
            0x80
        } else {
            0x00
        };
        q_opcd_a_t_r |= u8::from(self.op_code) << 3;
        q_opcd_a_t_r |= if self.authoritative { 0x4 } else { 0x0 };
        q_opcd_a_t_r |= if self.truncation { 0x2 } else { 0x0 };
        q_opcd_a_t_r |= if self.recursion_desired { 0x1 } else { 0x0 };
        encoder.emit(q_opcd_a_t_r)?;

        // IsRecursionAvailable, Triple 0's, ResponseCode
        let mut r_z_ad_cd_rcod: u8 = if self.recursion_available {
            0b1000_0000
        } else {
            0b0000_0000
        };
        r_z_ad_cd_rcod |= if self.authentic_data {
            0b0010_0000
        } else {
            0b0000_0000
        };
        r_z_ad_cd_rcod |= if self.checking_disabled {
            0b0001_0000
        } else {
            0b0000_0000
        };
        r_z_ad_cd_rcod |= self.response_code;
        encoder.emit(r_z_ad_cd_rcod)?;

        encoder.emit_u16(self.query_count)?;
        encoder.emit_u16(self.answer_count)?;
        encoder.emit_u16(self.name_server_count)?;
        encoder.emit_u16(self.additional_count)?;

        Ok(())
    }
}

impl<'r> BinDecodable<'r> for Header {
    fn read(decoder: &mut BinDecoder<'r>) -> ProtoResult<Self> {
        let id = decoder.read_u16()?.unverified(/*it is valid for this to be any u16*/);

        let q_opcd_a_t_r = decoder.pop()?.unverified(/*used as a bitfield, this is safe*/);
        // if the first bit is set
        let message_type = if (0b1000_0000 & q_opcd_a_t_r) == 0b1000_0000 {
            MessageType::Response
        } else {
            MessageType::Query
        };
        // the 4bit opcode, masked and then shifted right 3bits for the u8...
        let op_code: OpCode = OpCode::from_u8((0b0111_1000 & q_opcd_a_t_r) >> 3)?;
        let authoritative = (0b0000_0100 & q_opcd_a_t_r) == 0b0000_0100;
        let truncation = (0b0000_0010 & q_opcd_a_t_r) == 0b0000_0010;
        let recursion_desired = (0b0000_0001 & q_opcd_a_t_r) == 0b0000_0001;

        let r_z_ad_cd_rcod = decoder.pop()?.unverified(/*used as a bitfield, this is safe*/); // fail fast...

        let recursion_available = (0b1000_0000 & r_z_ad_cd_rcod) == 0b1000_0000;
        let authentic_data = (0b0010_0000 & r_z_ad_cd_rcod) == 0b0010_0000;
        let checking_disabled = (0b0001_0000 & r_z_ad_cd_rcod) == 0b0001_0000;
        let response_code: u8 = 0b0000_1111 & r_z_ad_cd_rcod;

        // TODO: We should pass these restrictions on, they can't be trusted, but that would seriosly complicate the Header type..
        // TODO: perhaps the read methods for BinDecodable should return Restrict?
        let query_count = decoder.read_u16()?.unverified(/*this must be verified when reading queries*/);
        let answer_count = decoder.read_u16()?.unverified(/*this must be evaluated when reading records*/);
        let name_server_count = decoder.read_u16()?.unverified(/*this must be evaluated when reading records*/);
        let additional_count = decoder.read_u16()?.unverified(/*this must be evaluated when reading records*/);

        // TODO: question, should this use the builder pattern instead? might be cleaner code, but
        //  this guarantees that the Header is fully instantiated with all values...
        Ok(Header {
            id,
            message_type,
            op_code,
            authoritative,
            truncation,
            recursion_desired,
            recursion_available,
            authentic_data,
            checking_disabled,
            response_code,
            query_count,
            answer_count,
            name_server_count,
            additional_count,
        })
    }
}

#[test]
fn test_parse() {
    let byte_vec = vec![
        0x01, 0x10, 0xAA, 0x83, // 0b1010 1010 1000 0011
        0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11,
    ];

    let mut decoder = BinDecoder::new(&byte_vec);

    let expect = Header {
        id: 0x0110,
        message_type: MessageType::Response,
        op_code: OpCode::Update,
        authoritative: false,
        truncation: true,
        recursion_desired: false,
        recursion_available: true,
        authentic_data: false,
        checking_disabled: false,
        response_code: ResponseCode::NXDomain.low(),
        query_count: 0x8877,
        answer_count: 0x6655,
        name_server_count: 0x4433,
        additional_count: 0x2211,
    };

    let got = Header::read(&mut decoder).unwrap();

    assert_eq!(got, expect);
}

#[test]
fn test_write() {
    let header = Header {
        id: 0x0110,
        message_type: MessageType::Response,
        op_code: OpCode::Update,
        authoritative: false,
        truncation: true,
        recursion_desired: false,
        recursion_available: true,
        authentic_data: false,
        checking_disabled: false,
        response_code: ResponseCode::NXDomain.low(),
        query_count: 0x8877,
        answer_count: 0x6655,
        name_server_count: 0x4433,
        additional_count: 0x2211,
    };

    let expect: Vec<u8> = vec![
        0x01, 0x10, 0xAA, 0x83, // 0b1010 1010 1000 0011
        0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11,
    ];

    let mut bytes = Vec::with_capacity(512);
    {
        let mut encoder = BinEncoder::new(&mut bytes);
        header.emit(&mut encoder).unwrap();
    }

    assert_eq!(bytes, expect);
}