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
//! Implementation of [coap_message::MutableWritableMessage] on a slice of memory
//!
//! [Message] is the main struct of this module.

pub use crate::error::WriteError;

/// A message writing into a preallocated buffer
pub struct Message<'a> {
    code: &'a mut u8,
    tail: &'a mut [u8],

    /// Latest option that has been written
    latest: u16,
    /// Index after tha last written byte
    end: usize,
    /// First byte of any written payload
    ///
    /// If this has been set, the byte before it was written 0xff. This is None if the was an empty
    /// payload.
    payload_start: Option<usize>,
}

/// A CoAP messag that resides in contiguous mutable memory
///
/// Exceeding the guarantees of MutableWritableMessage, this does allow some out-of-sequence
/// invocations: Even after payload has been written to, options can be added, memmove'ing
/// (right-rotating) the written payload, possibly truncating it out of the buffer. This is needed
/// to accommodate libOSCORE's requirements (because while libOSCORE can also do without a
/// memmove-ing message, that'd require its use through WritableMessage to adhere to in-OSCORE
/// write sequence conventions, making the whole situation no easier on the coap-message
/// abstraction). Data will only be moved if an option is added after content has been set, so this
/// comes at no runtime cost for those who do not need it. (It may be later turned into a feature.
/// Then, the memmove code would be removed; carrying the latest option number in the WriteState
/// should come at no extra cost due to the struct's content and alignment).
impl<'a> Message<'a> {
    pub fn new(code: &'a mut u8, tail: &'a mut [u8]) -> Self {
        Message {
            code,
            tail,
            latest: 0,
            end: 0,
            payload_start: None,
        }
    }

    /// Discard anything that has been written in the message, and allow any operation that would
    /// be allowed after calling [`Self::new()`].
    ///
    /// This is a short-cut to messages that can be snapshotted and rewound, and a band-aid until
    /// that is available in traits.
    pub fn reset(&mut self) {
        self.latest = 0;
        self.end = 0;
        self.payload_start = None;
        // This is not strictly necessary, but at least a temporarily useful thing that will keep
        // users from relying on the code to be persisted.
        *self.code = 0;
    }

    /// Create a MutableWritableMessage on a buffer that already contains a serialized message
    ///
    /// While this is generally not useful (a parsed message has its paylod already set, and if
    /// there's no payload, there's no space to add options or any payload), it allows mutable
    /// access to the option bytes and the payload. This is primarily useful in situations when
    /// data is processed in place, eg. decrypted (in OSCORE), or CBOR is shifted around to get
    /// contiguous slices out of indefinite length strings.
    ///
    /// This uses the same strategy for errors as [crate::inmemory::Message]: Validation happens
    /// incrementally.
    pub fn new_from_existing(code: &'a mut u8, tail: &'a mut [u8]) -> Self {
        // We'll use a read-only message to get the cursor parameters. This is not terribly great
        // performance-wise (this will iterate over the options, and then the user will do that
        // again); can still be improved to keep the cursor in an indefinite state until the
        // options have been iterated over, preferably when the inmemory message gains the
        // corresponding memoization.
        let read_only = crate::inmemory::Message::new(*code, tail);
        // We don't do thorough parsing, just enough to get an easy index. Either the message is
        // valid, then this will be OK, or it's not and the user will get garbage at read time.
        use coap_message::ReadableMessage;
        let payload_len = read_only.payload().len();

        let payload_start = match payload_len {
            0 => None,
            n => Some(tail.len() - n),
        };
        let end = tail.len();

        Message {
            code,
            tail,
            payload_start,
            end,
            // No point in determining what the actual last option is: This is for mutating
            // options, not adding new ones
            latest: u16::MAX,
        }
    }

    /// Return the number of bytes that were populated inside tail
    pub fn finish(self) -> usize {
        self.end
    }
}

impl<'a> coap_message::ReadableMessage for Message<'a> {
    type Code = u8;
    type MessageOption<'b> = crate::inmemory::MessageOption<'b>
    where
        Self: 'b,
    ;
    type OptionsIter<'b> = crate::inmemory::OptionsIter<'b>
    where
        Self: 'b,
    ;
    fn code(&self) -> u8 {
        *self.code
    }
    // Funny detail on the side: If this is called often, an inmemory_write::Message might be more
    // efficient even for plain reading than an inmemory::Message (because we don't have to iterate)
    fn payload(&self) -> &[u8] {
        match self.payload_start {
            None => &[],
            Some(start) => &self.tail[start..self.end],
        }
    }
    fn options(&self) -> <Self as coap_message::ReadableMessage>::OptionsIter<'_> {
        crate::inmemory::OptionsIter(crate::option_iteration::OptPayloadReader::new(self.tail))
    }
}

impl<'a> coap_message::MinimalWritableMessage for Message<'a> {
    type Code = u8;
    type OptionNumber = u16;
    type UnionError = WriteError;
    type AddOptionError = WriteError;
    type SetPayloadError = WriteError;

    fn set_code(&mut self, code: u8) {
        *self.code = code;
    }

    fn add_option(&mut self, number: u16, data: &[u8]) -> Result<(), WriteError> {
        let delta = number
            .checked_sub(self.latest)
            .ok_or(WriteError::OutOfSequence)?;
        self.latest = number;
        let encoded = crate::option_extension::encode_extensions(delta, data.len() as u16);
        let encoded = encoded.as_ref();
        let added_len = encoded.len() + data.len();

        let option_cursor;
        if let Some(payload_start) = self.payload_start {
            option_cursor = payload_start - 1;
            // Could also rotate_right, but we don't need the shifted-out bytes preserved in the
            // area we'll overwrite in the next instructions
            self.tail.copy_within(
                option_cursor..self.tail.len() - added_len,
                option_cursor + added_len,
            );
            self.payload_start = Some(payload_start + added_len);
        } else {
            option_cursor = self.end;
        }

        self.tail
            .get_mut(option_cursor..option_cursor + encoded.len())
            .ok_or(WriteError::OutOfSpace)?
            .copy_from_slice(encoded);
        let option_cursor = option_cursor + encoded.len();
        self.tail
            .get_mut(option_cursor..option_cursor + data.len())
            .ok_or(WriteError::OutOfSpace)?
            .copy_from_slice(data);
        self.end += added_len;

        Ok(())
    }

    fn set_payload(&mut self, payload: &[u8]) -> Result<(), WriteError> {
        if self.payload_start.is_some() {
            // We might allow double setting the payload through later extensions, but as for this
            // interface it's once only. We don't detect double setting of empty payloads, but it's
            // not this implementation's purpose to act as a linter.
            //
            // (And whoever uses the options-and-payload-mixed properties will use payload_mut
            // instead).
            return Err(WriteError::OutOfSequence);
        };
        if !payload.is_empty() {
            *self.tail.get_mut(self.end).ok_or(WriteError::OutOfSpace)? = 0xff;
            let start = self.end + 1;
            self.end = start + payload.len();
            self.tail
                .get_mut(start..self.end)
                .ok_or(WriteError::OutOfSpace)?
                .copy_from_slice(payload);
            self.payload_start = Some(start);
        }
        Ok(())
    }
}

impl<'a> coap_message::MutableWritableMessage for Message<'a> {
    fn available_space(&self) -> usize {
        self.tail.len()
            - self
                .payload_start
                .map(|s| s - 1) // available_space includes the payload indicator
                .unwrap_or(self.end)
    }

    fn payload_mut_with_len(&mut self, len: usize) -> Result<&mut [u8], WriteError> {
        if len == 0 {
            // Just finish the side effect and return something good enough; this allows the easier
            // path for the rest of the function to pick a start, end, and serve that.
            self.truncate(0)?;
            return Ok(&mut []);
        }

        let start = match self.payload_start {
            None => {
                self.tail[self.end] = 0xff;
                self.end + 1
            }
            Some(payload_start) => payload_start,
        };
        let end = start + len;

        let end = end.clamp(0, self.tail.len());

        self.payload_start = Some(start);
        self.end = end;
        self.tail.get_mut(start..end).ok_or(WriteError::OutOfSpace)
    }

    fn truncate(&mut self, len: usize) -> Result<(), WriteError> {
        match (len, self.payload_start) {
            (0, Some(payload_start)) => {
                self.end = payload_start - 1;
            }
            (0, None) => {}
            (len, Some(payload_start)) if self.end - payload_start >= len => {
                self.end = payload_start + len;
            }
            _ => return Err(WriteError::OutOfSpace),
        }
        Ok(())
    }

    fn mutate_options<F>(&mut self, mut f: F)
    where
        F: FnMut(u16, &mut [u8]),
    {
        // TBD this is excessively complex, and grounds for finding a better interface. ("Set
        // option and give me a key to update it later with a mutable reference")?

        let optend = self.payload_start.map(|s| s - 1).unwrap_or(self.end);

        // May end in a payload marker or just plain end
        let mut slice = &mut self.tail[..optend];

        let mut option_base = 0;

        while !slice.is_empty() {
            // This is copied and adapted from
            // coap_messsage_utils::option_iteration::OptPayloadReader and not used through it,
            // because that'd be a whole separate implementation there with mut.
            // (It's bad enough that take_extension needs the trickery)
            let delta_len = slice[0];
            slice = &mut slice[1..];

            if delta_len == 0xff {
                break;
            }

            let mut delta = (delta_len as u16) >> 4;
            let mut len = (delta_len as u16) & 0x0f;

            let new_len = {
                // Workaround for https://github.com/rust-lang/rust-clippy/issues/10608
                #[allow(clippy::redundant_slicing)]
                // To get take_extension to cooperate...
                let mut readable = &slice[..];

                crate::option_extension::take_extension(&mut delta, &mut readable)
                    .expect("Invalid encoded option in being-written message");
                crate::option_extension::take_extension(&mut len, &mut readable)
                    .expect("Invalid encoded option in being-written message");

                readable.len()
            };
            // ... and get back to a mutable form
            let trim = slice.len() - new_len;
            slice = &mut slice[trim..];

            option_base += delta;

            let len = len.into();
            f(option_base, &mut slice[..len]);
            slice = &mut slice[len..];
        }
    }
}