wasefire_protocol_usb/
common.rs

1// Copyright 2024 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Provides message fragmentation into packets.
16//!
17//! A packet is always 64 bytes. The first byte is a header:
18//! - Bit 0x80 is set for the first packet of a message (mutually exclusive with last)
19//! - Bit 0x40 is set for the last packet of a message (mutually exclusive with first)
20//! - Bit 0x01 is set when the fragment has a footer
21//! - Bits 0x3e are never set.
22//!
23//! The footer is the last byte of the packet. It contains the length of the content in bytes (a
24//! number between 0 and 62). Without a footer, the content length is 63 bytes.
25//!
26//! The content starts at the second byte (after the header). If it is shorter than 63 bytes (i.e.
27//! the packet has a footer), then the padding between the content and the footer is set to zero.
28//!
29//! A message always have at least 2 packets: the first packet and the last packet.
30
31use alloc::vec::Vec;
32
33use wasefire_logger as log;
34
35pub const PACKET_NO_REQUEST: [u8; 64] = [0; 64];
36
37#[derive(Debug, Copy, Clone, PartialEq, Eq)]
38pub struct Packet<'a> {
39    pub order: Order,
40    pub content: &'a [u8],
41}
42
43impl<'a> Packet<'a> {
44    pub fn decode(bytes: &'a [u8; 64]) -> Option<Self> {
45        let header = bytes[0];
46        ensure(header & 0x3e == 0)?;
47        let first = header & 0x80 != 0;
48        let last = header & 0x40 != 0;
49        let order = Order::new(first, last)?;
50        let mut content = &bytes[1 ..];
51        if header & 0x01 != 0 {
52            let length = bytes[63] as usize;
53            ensure((0 ..= 62).contains(&length))?;
54            ensure_zero(&content[length .. 62])?;
55            content = &content[.. length];
56        };
57        Some(Packet { order, content })
58    }
59
60    #[allow(clippy::identity_op)]
61    pub fn encode(self, bytes: &mut [u8; 64]) {
62        let Packet { order, content } = self;
63        let has_footer = content.len() != 63;
64        bytes[0] = 0;
65        bytes[0] |= 0x80 * matches!(order, Order::First) as u8;
66        bytes[0] |= 0x40 * matches!(order, Order::Last) as u8;
67        bytes[0] |= 0x01 * has_footer as u8;
68        bytes[1 ..][.. content.len()].copy_from_slice(content);
69        if has_footer {
70            bytes[1 + content.len() ..].fill(0);
71            bytes[63] = content.len() as u8;
72        }
73    }
74}
75
76#[derive(Debug, Copy, Clone, PartialEq, Eq)]
77pub enum Order {
78    First,
79    Middle,
80    Last,
81}
82
83impl Order {
84    fn new(first: bool, last: bool) -> Option<Self> {
85        Some(match (first, last) {
86            (true, true) => return None,
87            (true, false) => Order::First,
88            (false, true) => Order::Last,
89            (false, false) => Order::Middle,
90        })
91    }
92}
93
94#[derive(Default)]
95pub struct Encoder<'a> {
96    message: &'a [u8],
97    count: usize, // number of emitted fragments
98}
99
100impl<'a> Encoder<'a> {
101    pub fn new(message: &'a [u8]) -> Self {
102        Encoder { message, count: 0 }
103    }
104}
105
106impl Iterator for Encoder<'_> {
107    type Item = [u8; 64];
108
109    fn next(&mut self) -> Option<[u8; 64]> {
110        let total = core::cmp::max(2, self.message.len().div_ceil(63));
111        ensure(self.count < total)?;
112        let first = self.count == 0;
113        let start = core::cmp::min(self.message.len(), 63 * self.count);
114        let length = core::cmp::min(63, self.message.len().saturating_sub(63 * self.count));
115        let content = &self.message[start ..][.. length];
116        self.count += 1;
117        let last = self.count == total;
118        let order = Order::new(first, last).unwrap();
119        let mut packet = [0; 64];
120        Packet { order, content }.encode(&mut packet);
121        Some(packet)
122    }
123}
124
125#[derive(Default)]
126pub struct Decoder {
127    // None until a first packet is pushed.
128    message: Option<Vec<u8>>,
129}
130
131impl Decoder {
132    pub fn push(mut self, packet: &[u8; 64]) -> Option<Result<Vec<u8>, Decoder>> {
133        let Packet { order, content } = Packet::decode(packet)?;
134        if order == Order::First {
135            if self.message.is_some() {
136                log::warn!("Discarding previous message on repeated first packet.");
137            }
138            self.message = Some(content.to_vec());
139            return Some(Err(self));
140        }
141        let message = match &mut self.message {
142            Some(x) => x,
143            None => {
144                log::debug!("Discarding until first packet.");
145                return Some(Err(self));
146            }
147        };
148        message.extend_from_slice(content);
149        if order == Order::Last {
150            return Some(Ok(core::mem::take(message)));
151        }
152        Some(Err(self))
153    }
154}
155
156fn ensure_zero(xs: &[u8]) -> Option<()> {
157    xs.iter().all(|&x| x == 0).then_some(())
158}
159
160fn ensure(cond: bool) -> Option<()> {
161    cond.then_some(())
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    #[test]
169    fn message_at_least_two_packets() {
170        assert_eq!(Encoder::new(&[]).count(), 2);
171        assert_eq!(Encoder::new(&[0]).count(), 2);
172        assert_eq!(Encoder::new(&[0; 126]).count(), 2);
173        assert_eq!(Encoder::new(&[0; 127]).count(), 3);
174    }
175
176    #[test]
177    fn message_round_trip() {
178        let max_len = 3 * 63;
179        let pattern: Vec<u8> = (1 ..= max_len as u8).collect();
180        for len in 0 ..= max_len {
181            let expected = &pattern[.. len];
182            let mut decoder = Decoder::default();
183            let mut packets = Encoder::new(expected);
184            let result = loop {
185                let packet = packets.next().unwrap();
186                match decoder.push(&packet).unwrap() {
187                    Ok(x) => break x,
188                    Err(x) => decoder = x,
189                }
190            };
191            assert_eq!(packets.next(), None);
192            assert_eq!(result, expected);
193        }
194    }
195}
wasefire_protocol_usb/common.rs

wasefire_protocol_usb/
common.rs
