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
/// The [`CobsDecoder`] type is used to decode a stream of bytes to a
/// given mutable output slice. This is often useful when heap data
/// structures are not available, or when not all message bytes are
/// received at a single point in time.
#[derive(Debug)]
pub struct CobsDecoder<'a> {
    /// Destination slice for decoded message
    dest: &'a mut [u8],

    /// Index of next byte to write in `dest`
    dest_idx: usize,

    /// Decoder state as an enum
    state: DecoderState,
}

/// The [`DecoderState`] is used to track the current state of a
/// streaming decoder. This struct does not contain the output buffer
/// (or a reference to one), and can be used when streaming the decoded
/// output to a custom data type.
#[derive(Debug)]
pub enum DecoderState {
    /// State machine has not received any non-zero bytes
    Idle,

    /// 1-254 bytes, can be header or 00
    Grab(u8),

    /// 255 bytes, will be a header next
    GrabChain(u8),
}

fn add(to: &mut [u8], idx: usize, data: u8) -> Result<(), ()> {
    *to.get_mut(idx)
        .ok_or_else(|| ())? = data;
    Ok(())
}

/// [`DecodeResult`] represents the possible non-error outcomes of
/// pushing an encoded data byte into the [`DecoderState`] state machine
pub enum DecodeResult {
    /// The given input byte did not prompt an output byte, either because the
    /// state machine is still idle, or we have just processed a header byte.
    /// More data is needed to complete the message.
    NoData,

    /// We have received a complete and well-encoded COBS message. The
    /// contents of the associated output buffer may now be used
    DataComplete,

    /// The following byte should be appended to the current end of the decoded
    /// output buffer.
    /// More data is needed to complete the message.
    DataContinue(u8),
}

impl DecoderState {
    /// Push a single encoded byte into the state machine. If the input was
    /// unexpected, such as an early end of a framed message segment, an Error will
    /// be returned, and the current associated output buffer contents should be discarded.
    ///
    /// If a complete message is indicated, the decoding state machine will automatically
    /// reset itself to the Idle state, and may be used to begin decoding another message.
    ///
    /// NOTE: Sentinel value must be included in the input to this function for the
    /// decoding to complete
    pub fn feed(&mut self, data: u8) -> Result<DecodeResult, ()> {
        use DecoderState::*;
        use DecodeResult::*;
        let (ret, state) = match (&self, data) {
            // Currently Idle, received a terminator, ignore, stay idle
            (Idle, 0x00) => (Ok(NoData), Idle),

            // Currently Idle, received a byte indicating the
            // next 255 bytes have no zeroes, so we will have 254 unmodified
            // data bytes, then an overhead byte
            (Idle, 0xFF) => (Ok(NoData), GrabChain(0xFE)),

            // Currently Idle, received a byte indicating there will be a
            // zero that must be modified in the next 1..=254 bytes
            (Idle, n)    => (Ok(NoData), Grab(n - 1)),

            // We have reached the end of a data run indicated by an overhead
            // byte, AND we have recieved the message terminator. This was a
            // well framed message!
            (Grab(0), 0x00) => (Ok(DataComplete), Idle),

            // We have reached the end of a data run indicated by an overhead
            // byte, and the next segment of 254 bytes will have no modified
            // sentinel bytes
            (Grab(0), 0xFF) => {
                (Ok(DataContinue(0)), GrabChain(0xFE))
            },

            // We have reached the end of a data run indicated by an overhead
            // byte, and we will treat this byte as a modified sentinel byte.
            // place the sentinel byte in the output, and begin processing the
            // next non-sentinel sequence
            (Grab(0), n) => {
                (Ok(DataContinue(0)), Grab(n - 1))
            },

            // We were not expecting the sequence to terminate, but here we are.
            // Report an error due to early terminated message
            (Grab(_), 0) => {
                (Err(()), Idle)
            }

            // We have not yet reached the end of a data run, decrement the run
            // counter, and place the byte into the decoded output
            (Grab(i), n) =>  {
                (Ok(DataContinue(n)), Grab(*i - 1))
            },

            // We have reached the end of a data run indicated by an overhead
            // byte, AND we have recieved the message terminator. This was a
            // well framed message!
            (GrabChain(0), 0x00) => {
                (Ok(DataComplete), Idle)
            }

            // We have reached the end of a data run, and we will begin another
            // data run with an overhead byte expected at the end
            (GrabChain(0), 0xFF) => (Ok(NoData), GrabChain(0xFE)),

            // We have reached the end of a data run, and we will expect `n` data
            // bytes unmodified, followed by a sentinel byte that must be modified
            (GrabChain(0), n) => (Ok(NoData), Grab(n - 1)),

            // We were not expecting the sequence to terminate, but here we are.
            // Report an error due to early terminated message
            (GrabChain(_), 0) => {
                (Err(()), Idle)
            }

            // We have not yet reached the end of a data run, decrement the run
            // counter, and place the byte into the decoded output
            (GrabChain(i), n) => {
                (Ok(DataContinue(n)), GrabChain(*i - 1))
            },
        };

        *self = state;
        ret
    }
}

impl<'a> CobsDecoder<'a> {

    /// Create a new streaming Cobs Decoder. Provide the output buffer
    /// for the decoded message to be placed in
    pub fn new(dest: &'a mut [u8]) -> CobsDecoder<'a> {
        CobsDecoder {
            dest,
            dest_idx: 0,
            state: DecoderState::Idle,
        }
    }

    /// Push a single byte into the streaming CobsDecoder. Return values mean:
    ///
    /// * Ok(None) - State machine okay, more data needed
    /// * Ok(Some(N)) - A message of N bytes was successfully decoded
    /// * Err(M) - Message decoding failed, and M bytes were written to output
    ///
    /// NOTE: Sentinel value must be included in the input to this function for the
    /// decoding to complete
    pub fn feed(&mut self, data: u8) -> Result<Option<usize>, usize> {
        match self.state.feed(data) {
            Err(()) => Err(self.dest_idx),
            Ok(DecodeResult::NoData) => Ok(None),
            Ok(DecodeResult::DataContinue(n)) => {
                add(self.dest, self.dest_idx, n).map_err(|_| self.dest_idx)?;
                self.dest_idx += 1;
                Ok(None)
            }
            Ok(DecodeResult::DataComplete) => {
                Ok(Some(self.dest_idx))
            }
        }
    }

    /// Push a slice of bytes into the streaming CobsDecoder. Return values mean:
    ///
    /// * Ok(None) - State machine okay, more data needed
    /// * Ok(Some((N, M))) - A message of N bytes was successfully decoded,
    ///     using M bytes from `data` (and earlier data)
    /// * Err(J) - Message decoding failed, and J bytes were written to output
    ///
    /// NOTE: Sentinel value must be included in the input to this function for the
    /// decoding to complete
    pub fn push(&mut self, data: &[u8]) -> Result<Option<(usize, usize)>, usize> {
        for (consumed_idx, d) in data.iter().enumerate() {
            let x = self.feed(*d);
            if let Some(decoded_bytes_ct) = x? {
                // convert from index to number of bytes consumed
                return Ok(Some((decoded_bytes_ct, consumed_idx + 1)));
            }
        }

        Ok(None)
    }
}

// This needs to be a macro because `src` and `dst` could be the same or different.
macro_rules! decode_raw (
    ($src:ident, $dst:ident) => ({
        let mut source_index = 0;
        let mut dest_index = 0;

        // Stop at the first terminator, if any
        let src_end = if let Some(end) = $src.iter().position(|b| *b == 0) {
            end
        } else {
            $src.len()
        };

        while source_index < src_end {
            let code = $src[source_index];

            if source_index + code as usize > src_end && code != 1 {
                return Err(());
            }

            source_index += 1;

            // TODO: There are potential `panic!`s in these dest_index offsets
            for _ in 1..code {
                $dst[dest_index] = $src[source_index];
                source_index += 1;
                dest_index += 1;
            }

            if 0xFF != code && source_index < src_end {
                $dst[dest_index] = 0;
                dest_index += 1;
            }
        }

        DecodeReport {
            dst_used: dest_index,
            src_used: source_index,
        }
    })
);

/// Decodes the `source` buffer into the `dest` buffer.
///
/// This function uses the typical sentinel value of 0.
///
/// # Failures
///
/// This will return `Err(())` if there was a decoding error. Otherwise,
/// it will return `Ok(n)` where `n` is the length of the decoded message.
pub fn decode(source: &[u8], dest: &mut[u8]) -> Result<usize, ()> {
    let mut dec = CobsDecoder::new(dest);

    // Did we decode a message, using some or all of the buffer?
    match dec.push(source).or(Err(()))? {
        Some((d_used, _s_used)) => return Ok(d_used),
        None => {},
    }

    // If we consumed the entire buffer, but did NOT get a message,
    // AND the message did not end with a zero, try providing one to
    // complete the decoding.
    if source.last() != Some(&0) {
        // Explicitly push sentinel of zero
        if let Some((d_used, _s_used)) = dec.push(&[0]).or(Err(()))? {
            return Ok(d_used)
        }
    }

    // Nope, no early message, no missing terminator, just failed to decode
    Err(())
}

/// A report of the source and destination bytes used during in-place decoding
#[derive(Debug)]
pub struct DecodeReport {
    // The number of source bytes used, NOT INCLUDING the sentinel byte,
    // if there was one.
    pub src_used: usize,

    // The number of bytes of the source buffer that now include the
    // decoded result
    pub dst_used: usize,
}

/// Decodes a message in-place.
///
/// This is the same function as `decode_in_place`, but provides a report
/// of both the number of source bytes consumed as well as the size of the
/// destination used.
pub fn decode_in_place_report(buff: &mut[u8]) -> Result<DecodeReport, ()> {
    Ok(decode_raw!(buff, buff))
}

/// Decodes a message in-place.
///
/// This is the same function as `decode`, but replaces the encoded message
/// with the decoded message instead of writing to another buffer.
///
/// The returned `usize` is the number of bytes used for the DECODED value,
/// NOT the number of source bytes consumed during decoding.
pub fn decode_in_place(buff: &mut[u8]) -> Result<usize, ()> {
    Ok(decode_raw!(buff, buff).dst_used)
}

/// Decodes the `source` buffer into the `dest` buffer using an arbitrary sentinel value.
///
/// This is done by XOR-ing each byte of the source message with the chosen sentinel value,
/// which transforms the message into the same message encoded with a sentinel value of 0.
/// Then the regular decoding transformation is performed.
///
/// The returned `usize` is the number of bytes used for the DECODED value,
/// NOT the number of source bytes consumed during decoding.
pub fn decode_with_sentinel(source: &[u8], dest: &mut[u8], sentinel: u8) -> Result<usize, ()> {
    for (x, y) in source.iter().zip(dest.iter_mut()) {
        *y = *x ^ sentinel;
    }
    decode_in_place(dest)
}

/// Decodes a message in-place using an arbitrary sentinel value.
///
/// The returned `usize` is the number of bytes used for the DECODED value,
/// NOT the number of source bytes consumed during decoding.
pub fn decode_in_place_with_sentinel(buff: &mut[u8], sentinel: u8) -> Result<usize, ()> {
    for x in buff.iter_mut() {
        *x ^= sentinel;
    }
    decode_in_place(buff)
}

#[cfg(feature = "use_std")]
/// Decodes the `source` buffer into a vector.
pub fn decode_vec(source: &[u8]) -> Result<Vec<u8>, ()> {
    let mut decoded = vec![0; source.len()];
    match decode(source, &mut decoded[..]) {
        Ok(n) => {
            decoded.truncate(n);
            Ok(decoded)
        },
        Err(()) => Err(()),
    }
}

#[cfg(feature = "use_std")]
/// Decodes the `source` buffer into a vector with an arbitrary sentinel value.
pub fn decode_vec_with_sentinel(source: &[u8], sentinel: u8) -> Result<Vec<u8>, ()> {
    let mut decoded = vec![0; source.len()];
    match decode_with_sentinel(source, &mut decoded[..], sentinel) {
        Ok(n) => {
            decoded.truncate(n);
            Ok(decoded)
        },
        Err(()) => Err(()),
    }
}