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

//! This module provides [HeapMessage], an implementation of all the message traits backed by heap
//! memory.
//!
//! It serves both to easily get an owned copy of a message (using [`HeapMessage::new()`] and
//! [`MinimalWritableMessage::set_from_message()`]), especially if that needs frequent manipulation
//! of options, and as an example of how to the message traits can be implemented.
//!
//! This module is available only if the `alloc` feature is enabled on this crate.

use crate::message::*;

extern crate alloc;
use alloc::{vec, vec::Vec};

/// A heap CoAP message backed by allocated memory
///
/// It stores its payload in a [`Vec`], and uses a [`BTreeMap<_, Vec<_>>`](alloc::collections::BTreeMap) to
/// store all the individual option values.
///
/// It offers a few methods for direct manipulation of options, even out of sequence, that can not
/// be expected from a general message buffer and are thus not captured in traits.
///
/// ```
/// # use coap_message::{heapmessage::HeapMessage, ReadableMessage, MinimalWritableMessage};
/// let mut m = HeapMessage::new();
/// m.set_code(1);
/// m.add_option(11, b".well-known");
/// m.add_option(11, b"core");
///
/// let mut m2 = HeapMessage::new();
/// m2.set_from_message(&m);
/// assert!(m.code() == 1);
/// ```
#[derive(Debug)]
pub struct HeapMessage {
    code: u8,
    options: alloc::collections::BTreeMap<u16, Vec<Vec<u8>>>,
    payload: Vec<u8>,
}

impl HeapMessage {
    pub fn new() -> Self {
        Self {
            code: 0,
            options: alloc::collections::BTreeMap::new(),
            payload: vec![],
        }
    }

    // The following are not used via the WritableMessage as they're really custom to a
    // GUI-manager message

    /// Replace the occurrence'th value of the optnum option in the message
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn change_option(&mut self, optnum: u16, occurrence: usize, value: impl Into<Vec<u8>>) {
        self.options.get_mut(&optnum).unwrap()[occurrence] = value.into();
    }

    /// Remove the occurrence'th option of option number optnum
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn remove_option(&mut self, optnum: u16, occurrence: usize) {
        let vec = self.options.get_mut(&optnum).unwrap();
        vec.remove(occurrence);
        if vec.len() == 0 {
            self.options.remove(&optnum);
        }
    }

    /// Like MinimalWritableMessage::add_option, but allowing arbitrary access
    pub fn add_option(&mut self, optnum: u16, data: &[u8]) {
        self.options
            .entry(optnum)
            .or_insert(vec![])
            .push(data.to_vec());
    }
}

impl MinimalWritableMessage for HeapMessage {
    type Code = u8;
    type OptionNumber = u16;
    fn set_code(&mut self, code: u8) {
        self.code = code;
    }
    fn add_option(&mut self, optnum: u16, data: &[u8]) {
        // Only a debug-assert here because it's not something this needs to enforce, but more a
        // helper flag code that abuses an impl MinimalWritableMessage even when backed by the
        // HeapMessage that doesn't care
        debug_assert!(
            // Can be simplified once https://github.com/rust-lang/rust/issues/62924 is closed
            &optnum >= self.options.keys().max().unwrap_or(&0),
            "Option appended out-of-order"
        );

        self.add_option(optnum, data);
    }
    fn set_payload(&mut self, payload: &[u8]) {
        self.payload = payload.to_vec();
    }
}

pub struct MessageOption<'a> {
    number: u16,
    value: &'a [u8],
}

impl<'a> crate::message::MessageOption for MessageOption<'a> {
    fn number(&self) -> u16 {
        self.number
    }
    fn value(&self) -> &[u8] {
        self.value
    }
}

pub struct ReadCursor<'a> {
    iter: alloc::collections::btree_map::Iter<'a, u16, Vec<Vec<u8>>>,
    popped: (u16, &'a [Vec<u8>]),
}

impl<'a> Iterator for ReadCursor<'a> {
    type Item = MessageOption<'a>;
    fn next(&mut self) -> Option<MessageOption<'a>> {
        if self.popped.1.len() == 0 {
            self.popped = self.iter.next().map(|(k, v)| (*k, v.as_slice()))?;

            // Avoided by construction
            debug_assert!(
                self.popped.1.len() > 0,
                "HeapMessage had present but empty option"
            );
        }

        let (first, rest) = self.popped.1.split_at(1);
        self.popped.1 = rest;
        Some(MessageOption {
            number: self.popped.0,
            value: first[0].as_slice(),
        })
    }
}

impl WithSortedOptions for HeapMessage {
    // By virtue of the order of options used in BTreeMap
}

impl ReadableMessage for HeapMessage {
    type Code = u8;
    type MessageOption<'a> = MessageOption<'a>;
    type OptionsIter<'a> = ReadCursor<'a>;

    fn code(&self) -> u8 {
        self.code
    }
    fn payload(&self) -> &[u8] {
        &self.payload
    }
    fn options<'m>(&'m self) -> ReadCursor<'m> {
        ReadCursor {
            iter: self.options.iter(),
            popped: (0, &[]),
        }
    }
}