1#[cfg(test)]
2mod tests;
3
4use crate::{Error, Limits, Wire};
5use alloc::vec::Vec;
6use core::marker::PhantomData;
7use postcard::ser_flavors::Flavor;
8
9pub(crate) const HEADER: usize = 11;
10
11pub trait Method {
14 type Request: Wire;
17 type Response: Wire;
19 const ID: u64;
21 const NAME: &'static str;
23 const REQUEST_LIMITS: Limits;
25 const RESPONSE_LIMITS: Limits;
27}
28
29macro_rules! retained {
30 ($name:ident) => {
31 pub struct $name<M: Method> {
34 bytes: Vec<u8>,
35 method: PhantomData<fn() -> M>,
36 }
37 impl<M: Method> $name<M> {
38 pub fn as_bytes(&self) -> &[u8] {
40 &self.bytes
41 }
42 pub fn into_bytes(self) -> Vec<u8> {
44 self.bytes
45 }
46 }
47 };
48}
49retained!(RequestBytes);
50retained!(ResponseBytes);
51
52pub(crate) fn admit<T: Wire>(limits: Limits) -> Result<usize, Error> {
53 let shape = T::SHAPE;
54 if [shape.bytes, shape.heap, shape.nodes, shape.depth].contains(&usize::MAX) {
55 return Err(Error::SchemaLimit);
56 }
57 let size = shape.bytes.checked_add(HEADER).ok_or(Error::SchemaLimit)?;
58 let allocation = shape.heap.checked_add(size).ok_or(Error::SchemaLimit)?;
59 if size > limits.bytes
60 || allocation > limits.allocation
61 || shape.nodes > limits.nodes
62 || shape.depth > limits.depth
63 {
64 return Err(Error::SchemaLimit);
65 }
66 Ok(size)
67}
68
69pub(crate) struct Output {
72 pub(crate) bytes: Vec<u8>,
73 pub(crate) limit: usize,
74}
75impl Flavor for Output {
76 type Output = Vec<u8>;
77 fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
78 if self.bytes.len() == self.limit {
79 return Err(postcard::Error::SerializeBufferFull);
80 }
81 self.bytes.push(byte);
82 Ok(())
83 }
84 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
85 if bytes.len() > self.limit - self.bytes.len() {
86 return Err(postcard::Error::SerializeBufferFull);
87 }
88 self.bytes.extend_from_slice(bytes);
89 Ok(())
90 }
91 fn finalize(self) -> postcard::Result<Self::Output> {
92 Ok(self.bytes)
93 }
94}
95
96struct Size {
98 bytes: usize,
99 limit: usize,
100}
101impl Flavor for Size {
102 type Output = usize;
103 fn try_push(&mut self, _: u8) -> postcard::Result<()> {
104 self.try_extend(&[0])
105 }
106 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
107 if bytes.len() > self.limit - self.bytes {
108 return Err(postcard::Error::SerializeBufferFull);
109 }
110 self.bytes += bytes.len();
111 Ok(())
112 }
113 fn finalize(self) -> postcard::Result<usize> {
114 Ok(self.bytes)
115 }
116}
117
118fn encoded_len<T: Wire>(limits: Limits, value: &T) -> Result<usize, Error> {
119 let limit = admit::<T>(limits)?;
120 postcard::serialize_with_flavor(
121 value,
122 Size {
123 bytes: HEADER,
124 limit,
125 },
126 )
127 .map_err(|_| Error::Malformed)
128}
129
130pub(crate) struct Compare<'a> {
133 pub(crate) remaining: &'a [u8],
134 pub(crate) allowance: usize,
135 pub(crate) equal: bool,
136}
137impl Flavor for Compare<'_> {
138 type Output = bool;
139 #[inline(always)]
140 fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
141 if self.allowance == 0 {
142 return Err(postcard::Error::SerializeBufferFull);
143 }
144 self.allowance -= 1;
145 if let Some((first, rest)) = self.remaining.split_first() {
146 self.equal &= *first == byte;
147 self.remaining = rest;
148 } else {
149 self.equal = false;
150 }
151 Ok(())
152 }
153 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
154 if let [byte] = bytes {
155 return self.try_push(*byte);
156 }
157 if bytes.len() > self.allowance {
158 return Err(postcard::Error::SerializeBufferFull);
159 }
160 self.allowance -= bytes.len();
161 if let Some((prefix, rest)) = self.remaining.split_at_checked(bytes.len()) {
162 self.equal &= if cfg!(any(target_arch = "wasm32", test))
166 && core::ptr::eq(prefix.as_ptr(), bytes.as_ptr())
167 {
168 true
171 } else if cfg!(any(target_arch = "wasm32", test)) && bytes.len() >= 8 {
172 if bytes.len() >= 64 {
173 let (left, left_tail) = prefix.as_chunks::<16>();
174 let (right, right_tail) = bytes.as_chunks::<16>();
175 left.iter().zip(right).fold(0_u128, |difference, (a, b)| {
176 difference | (u128::from_ne_bytes(*a) ^ u128::from_ne_bytes(*b))
177 }) == 0
178 && left_tail == right_tail
179 } else {
180 let (left, left_tail) = prefix.as_chunks::<8>();
181 let (right, right_tail) = bytes.as_chunks::<8>();
182 left.iter().zip(right).fold(0_u64, |difference, (a, b)| {
183 difference | (u64::from_ne_bytes(*a) ^ u64::from_ne_bytes(*b))
184 }) == 0
185 && left_tail == right_tail
186 }
187 } else {
188 prefix == bytes
189 };
190 self.remaining = rest;
191 } else {
192 self.equal = false;
193 self.remaining = &[];
194 }
195 Ok(())
196 }
197 fn finalize(self) -> postcard::Result<Self::Output> {
198 Ok(self.equal && self.remaining.is_empty())
199 }
200}
201
202fn encode<T: Wire>(id: u64, direction: u8, limits: Limits, value: &T) -> Result<Vec<u8>, Error> {
203 let size = admit::<T>(limits)?;
204 let mut bytes = Vec::new();
205 bytes
206 .try_reserve_exact(size)
207 .map_err(|_| Error::Allocation)?;
208 bytes.extend_from_slice(&header(id, direction));
209 postcard::serialize_with_flavor(value, Output { bytes, limit: size })
210 .map_err(|_| Error::Malformed)
211}
212
213fn decode<'de, T: Wire + serde::Deserialize<'de>>(
214 id: u64,
215 direction: u8,
216 limits: Limits,
217 bytes: &'de [u8],
218) -> Result<T, Error> {
219 let size = admit::<T>(limits)?;
220 if bytes.len() > size || bytes.len() > limits.bytes {
221 return Err(Error::Limit);
222 }
223 let payload = payload(id, direction, bytes)?;
224 let (value, rest) = postcard::take_from_bytes::<T>(payload).map_err(|_| Error::Malformed)?;
225 if !rest.is_empty() {
226 return Err(Error::Trailing);
227 }
228 let equal = postcard::serialize_with_flavor(
229 &value,
230 Compare {
231 remaining: payload,
232 allowance: size - HEADER,
233 equal: true,
234 },
235 )
236 .map_err(|_| Error::Malformed)?;
237 if !equal {
238 return Err(Error::Noncanonical);
239 }
240 Ok(value)
241}
242
243pub fn encode_request<M: Method>(value: &M::Request) -> Result<RequestBytes<M>, Error> {
245 Ok(RequestBytes {
246 bytes: encode(M::ID, 0, M::REQUEST_LIMITS, value)?,
247 method: PhantomData,
248 })
249}
250pub fn encoded_request_len<M: Method>(value: &M::Request) -> Result<usize, Error> {
254 encoded_len(M::REQUEST_LIMITS, value)
255}
256pub fn encoded_response_len<M: Method>(value: &M::Response) -> Result<usize, Error> {
260 encoded_len(M::RESPONSE_LIMITS, value)
261}
262pub fn decode_request<'de, M: Method>(bytes: &'de [u8]) -> Result<M::Request, Error>
264where
265 M::Request: serde::Deserialize<'de>,
266{
267 decode(M::ID, 0, M::REQUEST_LIMITS, bytes)
268}
269pub fn encode_response<M: Method>(value: &M::Response) -> Result<ResponseBytes<M>, Error> {
271 Ok(ResponseBytes {
272 bytes: encode(M::ID, 1, M::RESPONSE_LIMITS, value)?,
273 method: PhantomData,
274 })
275}
276pub fn decode_response<'de, M: Method>(bytes: &'de [u8]) -> Result<M::Response, Error>
278where
279 M::Response: serde::Deserialize<'de>,
280{
281 decode(M::ID, 1, M::RESPONSE_LIMITS, bytes)
282}
283
284pub(crate) fn header(id: u64, direction: u8) -> [u8; HEADER] {
285 let mut header = [0; HEADER];
286 header[..2].copy_from_slice(b"IW");
287 header[2..10].copy_from_slice(&id.to_le_bytes());
288 header[10] = direction;
289 header
290}
291pub(crate) fn payload(id: u64, direction: u8, bytes: &[u8]) -> Result<&[u8], Error> {
292 if bytes.len() < HEADER {
293 return Err(Error::Malformed);
294 }
295 if &bytes[..2] != b"IW" || bytes[2..10] != id.to_le_bytes() || bytes[10] != direction {
296 return Err(Error::Identity);
297 }
298 Ok(&bytes[HEADER..])
299}