coap_message_implementations/
heap.rs

1//! This module provides [HeapMessage], an implementation of all the message traits backed by heap
2//! memory.
3//!
4//! It serves both to easily get an owned copy of a message (using [`HeapMessage::new()`] and
5//! [`MinimalWritableMessage::set_from_message()`]), especially if that needs frequent manipulation
6//! of options, and as an example of how to the message traits can be implemented.
7//!
8//! This module is available only if the `alloc` feature is enabled on this crate.
9//!
10//! Cross references:
11//!
12//! * The [coap-lite] crate's Packet struct is very similar to this, to the point where their
13//!   [MutableWritableMessage] implementations are identical. They differ by whether or not they
14//!   carry CoAP-over-UDP header information.
15//!
16//! * This implementation used to be part of the [coap_message] crate before its 0.3 version.
17//!
18//! [coap-lite]: https://crates.io/crates/coap-lite
19
20extern crate alloc;
21use alloc::{vec, vec::Vec};
22use core::convert::Infallible;
23
24use coap_message::{
25    MinimalWritableMessage, MutableWritableMessage, ReadableMessage, SeekWritableMessage,
26    WithSortedOptions,
27};
28
29/// A heap CoAP message backed by allocated memory
30///
31/// It stores its payload in a [`Vec`], and uses a [`BTreeMap<_, Vec<_>>`](alloc::collections::BTreeMap) to
32/// store all the individual option values.
33///
34/// It offers a few methods for direct manipulation of options, even out of sequence, that can not
35/// be expected from a general message buffer and are thus not captured in traits.
36///
37/// ```
38/// # use coap_message_implementations::heap::HeapMessage;
39/// use coap_message::{MinimalWritableMessage, ReadableMessage};
40/// let mut m = HeapMessage::new();
41/// m.set_code(1);
42/// m.add_option(11, b".well-known");
43/// m.add_option(11, b"core");
44///
45/// let mut m2 = HeapMessage::new();
46/// m2.set_from_message(&m);
47/// assert!(m.code() == 1);
48/// ```
49#[derive(Debug, Clone)]
50pub struct HeapMessage {
51    code: u8,
52    options: alloc::collections::BTreeMap<u16, Vec<Vec<u8>>>,
53    payload: Vec<u8>,
54}
55
56impl Default for HeapMessage {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl HeapMessage {
63    pub fn new() -> Self {
64        Self {
65            code: 0,
66            options: alloc::collections::BTreeMap::new(),
67            payload: vec![],
68        }
69    }
70
71    /// Replace the occurrence'th value of the optnum option in the message
72    ///
73    /// Panics if there's not ``occurrence + 1`` options of that number.
74    pub fn change_option(&mut self, optnum: u16, occurrence: usize, value: impl Into<Vec<u8>>) {
75        self.options.get_mut(&optnum).unwrap()[occurrence] = value.into();
76    }
77
78    /// Remove the occurrence'th option of option number optnum
79    ///
80    /// Panics if there's not ``occurrence + 1`` options of that number.
81    pub fn remove_option(&mut self, optnum: u16, occurrence: usize) {
82        let vec = self.options.get_mut(&optnum).unwrap();
83        vec.remove(occurrence);
84        if vec.is_empty() {
85            self.options.remove(&optnum);
86        }
87    }
88}
89
90impl MinimalWritableMessage for HeapMessage {
91    type AddOptionError = Infallible;
92    type SetPayloadError = Infallible;
93    type UnionError = Infallible;
94
95    type Code = u8;
96    type OptionNumber = u16;
97
98    fn set_code(&mut self, code: u8) {
99        self.code = code;
100    }
101
102    fn add_option(&mut self, optnum: u16, data: &[u8]) -> Result<(), Infallible> {
103        self.options.entry(optnum).or_default().push(data.to_vec());
104        Ok(())
105    }
106
107    fn set_payload(&mut self, payload: &[u8]) -> Result<(), Infallible> {
108        self.payload = payload.to_vec();
109        Ok(())
110    }
111}
112
113impl MutableWritableMessage for HeapMessage {
114    fn available_space(&self) -> usize {
115        usize::MAX
116    }
117
118    fn payload_mut_with_len(&mut self, len: usize) -> Result<&mut [u8], Infallible> {
119        self.payload.resize(len, 0);
120        Ok(&mut self.payload)
121    }
122
123    fn truncate(&mut self, len: usize) -> Result<(), Infallible> {
124        self.payload.truncate(len);
125        Ok(())
126    }
127
128    fn mutate_options<F>(&mut self, mut callback: F)
129    where
130        F: FnMut(Self::OptionNumber, &mut [u8]),
131    {
132        for (&number, ref mut values) in self.options.iter_mut() {
133            for v in values.iter_mut() {
134                callback(number, v);
135            }
136        }
137    }
138}
139
140impl SeekWritableMessage for HeapMessage {}
141
142pub struct MessageOption<'a> {
143    number: u16,
144    value: &'a [u8],
145}
146
147impl coap_message::MessageOption for MessageOption<'_> {
148    fn number(&self) -> u16 {
149        self.number
150    }
151    fn value(&self) -> &[u8] {
152        self.value
153    }
154}
155
156pub struct ReadCursor<'a> {
157    iter: alloc::collections::btree_map::Iter<'a, u16, Vec<Vec<u8>>>,
158    popped: (u16, &'a [Vec<u8>]),
159}
160
161impl<'a> Iterator for ReadCursor<'a> {
162    type Item = MessageOption<'a>;
163    fn next(&mut self) -> Option<MessageOption<'a>> {
164        if self.popped.1.is_empty() {
165            self.popped = self.iter.next().map(|(k, v)| (*k, v.as_slice()))?;
166
167            // Avoided by construction
168            debug_assert!(
169                !self.popped.1.is_empty(),
170                "HeapMessage had present but empty option"
171            );
172        }
173
174        let (first, rest) = self.popped.1.split_at(1);
175        self.popped.1 = rest;
176        Some(MessageOption {
177            number: self.popped.0,
178            value: first[0].as_slice(),
179        })
180    }
181}
182
183impl WithSortedOptions for HeapMessage {
184    // By virtue of the order of options used in BTreeMap
185}
186
187impl ReadableMessage for HeapMessage {
188    type Code = u8;
189    type MessageOption<'a> = MessageOption<'a>;
190    type OptionsIter<'a> = ReadCursor<'a>;
191
192    fn code(&self) -> u8 {
193        self.code
194    }
195    fn payload(&self) -> &[u8] {
196        &self.payload
197    }
198    fn options(&self) -> ReadCursor<'_> {
199        ReadCursor {
200            iter: self.options.iter(),
201            popped: (0, &[]),
202        }
203    }
204}