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
//! Encodable trait & Encoder
use crate::error::{EncodeError, EncodeResult};

/// A trait for types which are deserializable to DHCP binary formats
pub trait Encodable {
    /// Read the type from the stream
    fn encode(&self, e: &mut Encoder<'_>) -> EncodeResult<()>;

    /// encode this type into a new `Vec`
    fn to_vec(&self) -> EncodeResult<Vec<u8>> {
        let mut buffer = Vec::with_capacity(512);
        let mut encoder = Encoder::new(&mut buffer);
        self.encode(&mut encoder)?;
        Ok(buffer)
    }
}
/// Encoder type, holds a mut ref to a buffer
/// that it will write data to and an offset
/// of the next position to write
#[derive(Debug)]
pub struct Encoder<'a> {
    buffer: &'a mut Vec<u8>,
    offset: usize,
}

impl<'a> Encoder<'a> {
    /// Create a new Encoder from a mutable buffer
    pub fn new(buffer: &'a mut Vec<u8>) -> Self {
        Self { buffer, offset: 0 }
    }

    /// Get a reference to the underlying buffer
    pub fn buffer(&self) -> &[u8] {
        self.buffer
    }

    /// write bytes to buffer
    /// Return:
    ///     number of bytes written
    pub fn write_slice(&mut self, bytes: &[u8]) -> EncodeResult<()> {
        let additional = bytes.len();
        // space already reserved, we may not need this
        if self.offset + additional <= self.buffer.len() {
            // if self.offset == self.buffer.len() indexing can panic
            for (byte, b) in self.buffer[self.offset..].iter_mut().zip(bytes.iter()) {
                *byte = *b;
            }
        } else {
            let expected_len = self.buffer.len() + additional;
            self.buffer.reserve(additional);
            self.buffer.extend_from_slice(bytes);

            debug_assert!(self.buffer.len() == expected_len);
        }

        let index = self
            .offset
            .checked_add(additional)
            .ok_or(EncodeError::AddOverflow)?;
        self.offset = index;
        Ok(())
    }

    /// Write const number of bytes to buffer
    pub fn write<const N: usize>(&mut self, bytes: [u8; N]) -> EncodeResult<()> {
        // TODO: refactor this and above method?
        // only difference is zip & extend
        let additional = bytes.len();
        // space already reserved, we may not need this
        if self.offset + additional <= self.buffer.len() {
            // if self.offset == self.buffer.len() indexing can panic
            for (byte, b) in self.buffer[self.offset..].iter_mut().zip(bytes) {
                *byte = b;
            }
        } else {
            let expected_len = self.buffer.len() + additional;
            self.buffer.reserve(additional);
            self.buffer.extend(bytes);
            debug_assert!(self.buffer.len() == expected_len);
        }

        let index = self
            .offset
            .checked_add(additional)
            .ok_or(EncodeError::AddOverflow)?;
        self.offset = index;
        Ok(())
    }

    /// write a u8
    pub fn write_u8(&mut self, data: u8) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// write a u16
    pub fn write_u16(&mut self, data: u16) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// write a u32
    pub fn write_u32(&mut self, data: u32) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// write a u128
    pub fn write_u128(&mut self, data: u128) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// write a u64
    pub fn write_u64(&mut self, data: u64) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// write a i32
    pub fn write_i32(&mut self, data: i32) -> EncodeResult<()> {
        self.write(data.to_be_bytes())
    }
    /// Writes bytes to buffer and pads with 0 bytes up to some fill_len
    ///
    /// Returns
    ///    Err - if bytes.len() is greater then fill_len
    pub fn write_fill_bytes(&mut self, bytes: &[u8], fill_len: usize) -> EncodeResult<()> {
        if bytes.len() > fill_len {
            return Err(EncodeError::StringSizeTooBig { len: bytes.len() });
        }
        let nul_len = fill_len - bytes.len();
        self.write_slice(bytes)?;
        for _ in 0..nul_len {
            self.write_u8(0)?;
        }
        Ok(())
    }
    /// Writes string to buffer and pads with 0 bytes up to some fill_len
    /// if String is None then write fill_len 0 bytes
    ///
    /// Returns
    ///    Err - if bytes.len() is greater then fill_len
    pub fn write_fill_string(&mut self, s: &Option<String>, fill_len: usize) -> EncodeResult<()> {
        match s {
            Some(sname) => {
                let bytes = sname.as_bytes();
                self.write_fill_bytes(bytes, fill_len)?;
            }
            None => {
                // should we keep some static [0;64] arrays around
                // to fill quickly?
                for _ in 0..fill_len {
                    self.write_u8(0)?;
                }
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn basic_encode() -> EncodeResult<()> {
        let mut buf = vec![0, 1, 2, 3, 4, 5];
        let mut enc = Encoder::new(&mut buf);
        enc.offset = 4;
        // write already reserved space
        enc.write_slice(&[5, 6])?;
        assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3, 5, 6]);
        assert_eq!(enc.offset, 6);
        // reserve extra space
        enc.write_slice(&[7, 8])?;
        assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3, 5, 6, 7, 8]);
        assert_eq!(enc.offset, 8);

        // start w/ empty buf
        let mut buf = vec![];
        let mut enc = Encoder::new(&mut buf);
        // reserve space & write
        enc.write_slice(&[0, 1, 2, 3])?;
        assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3]);
        assert_eq!(enc.offset, 4);
        Ok(())
    }
}