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;
16 type Response: Wire;
18 const ID: u64;
20 const NAME: &'static str;
22 const REQUEST_LIMITS: Limits;
24 const RESPONSE_LIMITS: Limits;
26}
27
28macro_rules! retained {
29 ($name:ident) => {
30 pub struct $name<M: Method> {
33 bytes: Vec<u8>,
34 method: PhantomData<fn() -> M>,
35 }
36 impl<M: Method> $name<M> {
37 pub fn as_bytes(&self) -> &[u8] {
39 &self.bytes
40 }
41 pub fn into_bytes(self) -> Vec<u8> {
43 self.bytes
44 }
45 }
46 };
47}
48retained!(RequestBytes);
49retained!(ResponseBytes);
50
51pub(crate) fn admit<T: Wire>(limits: Limits) -> Result<usize, Error> {
52 let shape = T::SHAPE;
53 if [shape.bytes, shape.heap, shape.nodes, shape.depth].contains(&usize::MAX) {
54 return Err(Error::SchemaLimit);
55 }
56 let size = shape.bytes.checked_add(HEADER).ok_or(Error::SchemaLimit)?;
57 let allocation = shape.heap.checked_add(size).ok_or(Error::SchemaLimit)?;
58 if size > limits.bytes
59 || allocation > limits.allocation
60 || shape.nodes > limits.nodes
61 || shape.depth > limits.depth
62 {
63 return Err(Error::SchemaLimit);
64 }
65 Ok(size)
66}
67
68pub(crate) struct Output {
71 pub(crate) bytes: Vec<u8>,
72 pub(crate) limit: usize,
73}
74impl Flavor for Output {
75 type Output = Vec<u8>;
76 fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
77 if self.bytes.len() == self.limit {
78 return Err(postcard::Error::SerializeBufferFull);
79 }
80 self.bytes.push(byte);
81 Ok(())
82 }
83 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
84 if bytes.len() > self.limit - self.bytes.len() {
85 return Err(postcard::Error::SerializeBufferFull);
86 }
87 self.bytes.extend_from_slice(bytes);
88 Ok(())
89 }
90 fn finalize(self) -> postcard::Result<Self::Output> {
91 Ok(self.bytes)
92 }
93}
94
95struct Size {
97 bytes: usize,
98 limit: usize,
99}
100impl Flavor for Size {
101 type Output = usize;
102 fn try_push(&mut self, _: u8) -> postcard::Result<()> {
103 self.try_extend(&[0])
104 }
105 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
106 if bytes.len() > self.limit - self.bytes {
107 return Err(postcard::Error::SerializeBufferFull);
108 }
109 self.bytes += bytes.len();
110 Ok(())
111 }
112 fn finalize(self) -> postcard::Result<usize> {
113 Ok(self.bytes)
114 }
115}
116
117fn encoded_len<T: Wire>(limits: Limits, value: &T) -> Result<usize, Error> {
118 let limit = admit::<T>(limits)?;
119 postcard::serialize_with_flavor(
120 value,
121 Size {
122 bytes: HEADER,
123 limit,
124 },
125 )
126 .map_err(|_| Error::Malformed)
127}
128
129pub(crate) struct Compare<'a> {
132 pub(crate) remaining: &'a [u8],
133 pub(crate) allowance: usize,
134 pub(crate) equal: bool,
135}
136impl Flavor for Compare<'_> {
137 type Output = bool;
138 #[inline(always)]
139 fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
140 if self.allowance == 0 {
141 return Err(postcard::Error::SerializeBufferFull);
142 }
143 self.allowance -= 1;
144 if let Some((first, rest)) = self.remaining.split_first() {
145 self.equal &= *first == byte;
146 self.remaining = rest;
147 } else {
148 self.equal = false;
149 }
150 Ok(())
151 }
152 fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
153 if let [byte] = bytes {
154 return self.try_push(*byte);
155 }
156 if bytes.len() > self.allowance {
157 return Err(postcard::Error::SerializeBufferFull);
158 }
159 self.allowance -= bytes.len();
160 if let Some((prefix, rest)) = self.remaining.split_at_checked(bytes.len()) {
161 self.equal &= if cfg!(any(target_arch = "wasm32", test)) && bytes.len() >= 8 {
165 let (left, left_tail) = prefix.as_chunks::<8>();
166 let (right, right_tail) = bytes.as_chunks::<8>();
167 left.iter()
168 .zip(right)
169 .all(|(a, b)| u64::from_ne_bytes(*a) == u64::from_ne_bytes(*b))
170 && left_tail == right_tail
171 } else {
172 prefix == bytes
173 };
174 self.remaining = rest;
175 } else {
176 self.equal = false;
177 self.remaining = &[];
178 }
179 Ok(())
180 }
181 fn finalize(self) -> postcard::Result<Self::Output> {
182 Ok(self.equal && self.remaining.is_empty())
183 }
184}
185
186fn encode<T: Wire>(id: u64, direction: u8, limits: Limits, value: &T) -> Result<Vec<u8>, Error> {
187 let size = admit::<T>(limits)?;
188 let mut bytes = Vec::new();
189 bytes
190 .try_reserve_exact(size)
191 .map_err(|_| Error::Allocation)?;
192 bytes.extend_from_slice(&header(id, direction));
193 postcard::serialize_with_flavor(value, Output { bytes, limit: size })
194 .map_err(|_| Error::Malformed)
195}
196
197fn decode<T: Wire>(id: u64, direction: u8, limits: Limits, bytes: &[u8]) -> Result<T, Error> {
198 let size = admit::<T>(limits)?;
199 if bytes.len() > size || bytes.len() > limits.bytes {
200 return Err(Error::Limit);
201 }
202 let payload = payload(id, direction, bytes)?;
203 let (value, rest) = postcard::take_from_bytes::<T>(payload).map_err(|_| Error::Malformed)?;
204 if !rest.is_empty() {
205 return Err(Error::Trailing);
206 }
207 let equal = postcard::serialize_with_flavor(
208 &value,
209 Compare {
210 remaining: payload,
211 allowance: size - HEADER,
212 equal: true,
213 },
214 )
215 .map_err(|_| Error::Malformed)?;
216 if !equal {
217 return Err(Error::Noncanonical);
218 }
219 Ok(value)
220}
221
222pub fn encode_request<M: Method>(value: &M::Request) -> Result<RequestBytes<M>, Error> {
224 Ok(RequestBytes {
225 bytes: encode(M::ID, 0, M::REQUEST_LIMITS, value)?,
226 method: PhantomData,
227 })
228}
229pub fn encoded_request_len<M: Method>(value: &M::Request) -> Result<usize, Error> {
233 encoded_len(M::REQUEST_LIMITS, value)
234}
235pub fn encoded_response_len<M: Method>(value: &M::Response) -> Result<usize, Error> {
239 encoded_len(M::RESPONSE_LIMITS, value)
240}
241pub fn decode_request<M: Method>(bytes: &[u8]) -> Result<M::Request, Error> {
243 decode(M::ID, 0, M::REQUEST_LIMITS, bytes)
244}
245pub fn encode_response<M: Method>(value: &M::Response) -> Result<ResponseBytes<M>, Error> {
247 Ok(ResponseBytes {
248 bytes: encode(M::ID, 1, M::RESPONSE_LIMITS, value)?,
249 method: PhantomData,
250 })
251}
252pub fn decode_response<M: Method>(bytes: &[u8]) -> Result<M::Response, Error> {
254 decode(M::ID, 1, M::RESPONSE_LIMITS, bytes)
255}
256
257pub(crate) fn header(id: u64, direction: u8) -> [u8; HEADER] {
258 let mut header = [0; HEADER];
259 header[..2].copy_from_slice(b"IW");
260 header[2..10].copy_from_slice(&id.to_le_bytes());
261 header[10] = direction;
262 header
263}
264pub(crate) fn payload(id: u64, direction: u8, bytes: &[u8]) -> Result<&[u8], Error> {
265 if bytes.len() < HEADER {
266 return Err(Error::Malformed);
267 }
268 if &bytes[..2] != b"IW" || bytes[2..10] != id.to_le_bytes() || bytes[10] != direction {
269 return Err(Error::Identity);
270 }
271 Ok(&bytes[HEADER..])
272}