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
/*
 * 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.
 */

use crate::serialize::binary::Restrict;
use thiserror::Error;

/// This is non-destructive to the inner buffer, b/c for pointer types we need to perform a reverse
///  seek to lookup names
///
/// A note on serialization, there was a thought to have this implement the Serde deserializer,
///  but given that this is such a small subset of all the serialization which that performs
///  this is a simpler implementation without the cruft, at least for serializing to/from the
///  binary DNS protocols.
pub struct BinDecoder<'a> {
    buffer: &'a [u8],    // The entire original buffer
    remaining: &'a [u8], // The unread section of the original buffer, so that reads do not cause a bounds check at the current seek offset
}

pub(crate) type DecodeResult<T> = Result<T, DecodeError>;

/// An error that can occur deep in a decoder
/// This type is kept very small so that function that use it inline often
#[derive(Clone, Copy, Debug, Error)]
#[non_exhaustive]
pub enum DecodeError {
    /// Insufficient data in the buffer for a read operation
    #[error("unexpected end of input reached")]
    InsufficientBytes,

    /// slice_from was called with an invalid index
    #[error(
        "the index passed to BinDecoder::slice_from must be greater than the decoder position"
    )]
    InvalidPreviousIndex,

    /// Pointer points to an index within or after the current label
    #[error("label points to data not prior to idx: {idx} ptr: {ptr}")]
    PointerNotPriorToLabel {
        /// index of the label containing this pointer
        idx: usize,
        /// location to which the pointer is directing
        ptr: u16,
    },

    /// Label bytes exceeded the limit of 63
    #[error("label bytes exceed 63: {0}")]
    LabelBytesTooLong(usize),

    /// An unrecognized label code was found
    #[error("unrecognized label code: {0:b}")]
    UnrecognizedLabelCode(u8),

    /// A domain name was too long
    #[error("name label data exceed 255: {0}")]
    DomainNameTooLong(usize),

    /// Overlapping labels
    #[error("overlapping labels name {label} other {other}")]
    LabelOverlapsWithOther {
        /// Start of the label that is overlaps
        label: usize,
        /// Start of the other label
        other: usize,
    },
}

impl<'a> BinDecoder<'a> {
    /// Creates a new BinDecoder
    ///
    /// # Arguments
    ///
    /// * `buffer` - buffer from which all data will be read
    pub fn new(buffer: &'a [u8]) -> Self {
        BinDecoder {
            buffer,
            remaining: buffer,
        }
    }

    /// Pop one byte from the buffer
    pub fn pop(&mut self) -> DecodeResult<Restrict<u8>> {
        if let Some((first, remaining)) = self.remaining.split_first() {
            self.remaining = remaining;
            return Ok(Restrict::new(*first));
        }
        Err(DecodeError::InsufficientBytes)
    }

    /// Returns the number of bytes in the buffer
    ///
    /// ```
    /// use trust_dns_proto::serialize::binary::BinDecoder;
    ///
    /// let deadbeef = b"deadbeef";
    /// let mut decoder = BinDecoder::new(deadbeef);
    /// assert_eq!(decoder.len(), 8);
    /// decoder.read_slice(7).unwrap();
    /// assert_eq!(decoder.len(), 1);
    /// ```
    pub fn len(&self) -> usize {
        self.remaining.len()
    }

    /// Returns `true` if the buffer is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Peed one byte forward, without moving the current index forward
    pub fn peek(&self) -> Option<Restrict<u8>> {
        Some(Restrict::new(*self.remaining.get(0)?))
    }

    /// Returns the current index in the buffer
    pub fn index(&self) -> usize {
        self.buffer.len() - self.remaining.len()
    }

    /// This is a pretty efficient clone, as the buffer is never cloned, and only the index is set
    ///  to the value passed in
    pub fn clone(&self, index_at: u16) -> BinDecoder<'a> {
        BinDecoder {
            buffer: self.buffer,
            remaining: &self.buffer[index_at as usize..],
        }
    }

    /// Reads a String from the buffer
    ///
    /// ```text
    /// <character-string> is a single
    /// length octet followed by that number of characters.  <character-string>
    /// is treated as binary information, and can be up to 256 characters in
    /// length (including the length octet).
    /// ```
    ///
    /// # Returns
    ///
    /// A String version of the character data
    pub fn read_character_data(&mut self) -> DecodeResult<Restrict<&[u8]>> {
        let length = self.pop()?.unverified() as usize;
        self.read_slice(length)
    }

    /// Reads a Vec out of the buffer
    ///
    /// # Arguments
    ///
    /// * `len` - number of bytes to read from the buffer
    ///
    /// # Returns
    ///
    /// The Vec of the specified length, otherwise an error
    pub fn read_vec(&mut self, len: usize) -> DecodeResult<Restrict<Vec<u8>>> {
        self.read_slice(len).map(|s| s.map(ToOwned::to_owned))
    }

    /// Reads a slice out of the buffer, without allocating
    ///
    /// # Arguments
    ///
    /// * `len` - number of bytes to read from the buffer
    ///
    /// # Returns
    ///
    /// The slice of the specified length, otherwise an error
    pub fn read_slice(&mut self, len: usize) -> DecodeResult<Restrict<&'a [u8]>> {
        if len > self.remaining.len() {
            return Err(DecodeError::InsufficientBytes);
        }
        let (read, remaining) = self.remaining.split_at(len);
        self.remaining = remaining;
        Ok(Restrict::new(read))
    }

    /// Reads a slice from a previous index to the current
    pub fn slice_from(&self, index: usize) -> DecodeResult<&'a [u8]> {
        if index > self.index() {
            return Err(DecodeError::InvalidPreviousIndex);
        }

        Ok(&self.buffer[index..self.index()])
    }

    /// Reads a byte from the buffer, equivalent to `Self::pop()`
    pub fn read_u8(&mut self) -> DecodeResult<Restrict<u8>> {
        self.pop()
    }

    /// Reads the next 2 bytes into u16
    ///
    /// This performs a byte-by-byte manipulation, there
    ///  which means endianness is implicitly handled (i.e. no network to little endian (intel), issues)
    ///
    /// # Return
    ///
    /// Return the u16 from the buffer
    pub fn read_u16(&mut self) -> DecodeResult<Restrict<u16>> {
        Ok(self
            .read_slice(2)?
            .map(|s| u16::from_be_bytes([s[0], s[1]])))
    }

    /// Reads the next four bytes into i32.
    ///
    /// This performs a byte-by-byte manipulation, there
    ///  which means endianness is implicitly handled (i.e. no network to little endian (intel), issues)
    ///
    /// # Return
    ///
    /// Return the i32 from the buffer
    pub fn read_i32(&mut self) -> DecodeResult<Restrict<i32>> {
        Ok(self.read_slice(4)?.map(|s| {
            assert!(s.len() == 4);
            i32::from_be_bytes([s[0], s[1], s[2], s[3]])
        }))
    }

    /// Reads the next four bytes into u32.
    ///
    /// This performs a byte-by-byte manipulation, there
    ///  which means endianness is implicitly handled (i.e. no network to little endian (intel), issues)
    ///
    /// # Return
    ///
    /// Return the u32 from the buffer
    pub fn read_u32(&mut self) -> DecodeResult<Restrict<u32>> {
        Ok(self.read_slice(4)?.map(|s| {
            assert!(s.len() == 4);
            u32::from_be_bytes([s[0], s[1], s[2], s[3]])
        }))
    }
}

#[cfg(tests)]
mod tests {
    use super::*;

    #[test]
    fn test_read_slice() {
        let deadbeef = b"deadbeef";
        let mut decoder = BinDecoder::new(deadbeef);

        let read = decoder.read_slice(4).expect("failed to read dead");
        assert_eq!(read, "dead");

        let read = decoder.read_slice(2).expect("failed to read be");
        assert_eq!(read, "be");

        let read = decoder.read_slice(0).expect("failed to read nothing");
        assert_eq!(read, "");

        // this should fail
        assert!(decoder.read_slice(3).is_err());
    }

    #[test]
    fn test_read_slice_from() {
        let deadbeef = b"deadbeef";
        let mut decoder = BinDecoder::new(deadbeef);

        decoder.read_slice_from(4).expect("failed to read dead");
        let read = decoder.slice_from(0).expect("failed to get slice");
        assert_eq!(read, "dead");

        decoder.read_slice(2).expect("failed to read be");
        let read = decoder.slice_from(4).expect("failed to get slice");
        assert_eq!(read, "be");

        decoder.read_slice(0).expect("failed to read nothing");
        let read = decoder.slice_from(4).expect("failed to get slice");
        assert_eq!(read, "be");

        // this should fail
        assert!(decoder.slice_from(6).is_err());
        assert!(decoder.slice_from(10).is_err());
    }
}