ic_wire/contract/
mod.rs

1#[cfg(test)]
2mod tests;
3
4use crate::{Error, Limits, Wire};
5use alloc::vec::Vec;
6use core::{cell::Cell, marker::PhantomData};
7use postcard::ser_flavors::Flavor;
8
9pub(crate) const HEADER: usize = 11;
10
11/// One consumer-owned endpoint contract. IDs must be unique across incompatible
12/// layouts and methods. A source rename alone must not change its identity.
13pub trait Method {
14    /// Typed request, including any consumer-owned proof/envelope.
15    /// Borrowed fields can use lifetimes on the implementing method marker.
16    type Request: Wire;
17    /// Typed response, typically `Result<Success, ApplicationError>`.
18    type Response: Wire;
19    /// Stable contract identity assigned by the consumer.
20    const ID: u64;
21    /// Exact exported IC method name; used by typed transport adapters.
22    const NAME: &'static str;
23    /// Request admission budget.
24    const REQUEST_LIMITS: Limits;
25    /// Reply admission budget.
26    const RESPONSE_LIMITS: Limits;
27}
28
29macro_rules! retained {
30    ($name:ident) => {
31        /// Validated, immutable bytes bound to a method and direction. Retain
32        /// these for retries; the codec does not decide when retry is safe.
33        pub struct $name<M: Method> {
34            bytes: Vec<u8>,
35            method: PhantomData<fn() -> M>,
36        }
37        impl<M: Method> $name<M> {
38            /// Borrow the exact encoding for transport or hashing.
39            pub fn as_bytes(&self) -> &[u8] {
40                &self.bytes
41            }
42            /// Transfer ownership to a transport or persistence owner.
43            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
69// Reserve fallibly once, then initialize only emitted bytes. Capacity is not
70// the schema limit: allocators may return more capacity than requested.
71pub(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
96// Finite encoding visits the value once. Grow only for emitted bytes, keeping
97// allocation requests within the schema ceiling even when Vec has spare capacity.
98struct GrowingOutput<'a> {
99    output: Output,
100    allocation_failed: &'a Cell<bool>,
101}
102impl GrowingOutput<'_> {
103    #[cold]
104    #[inline(never)]
105    fn reserve(&mut self, additional: usize) -> postcard::Result<()> {
106        let bytes = &mut self.output.bytes;
107        if additional > self.output.limit - bytes.len() {
108            return Err(postcard::Error::SerializeBufferFull);
109        }
110        let required = bytes.len() + additional;
111        if required > bytes.capacity() {
112            // Once a message is substantial relative to its ceiling, reserve
113            // the remainder to avoid repeated copies of large batch payloads.
114            let capacity = if required > self.output.limit / 4 {
115                self.output.limit
116            } else {
117                required
118                    .max(bytes.capacity().saturating_mul(2).max(64))
119                    .min(self.output.limit)
120            };
121            bytes
122                .try_reserve_exact(capacity - bytes.len())
123                .map_err(|_| {
124                    self.allocation_failed.set(true);
125                    postcard::Error::SerializeBufferFull
126                })?;
127        }
128        Ok(())
129    }
130}
131impl Flavor for GrowingOutput<'_> {
132    type Output = Vec<u8>;
133    #[inline(always)]
134    fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
135        if self.output.bytes.len() == self.output.limit {
136            return Err(postcard::Error::SerializeBufferFull);
137        }
138        if self.output.bytes.len() == self.output.bytes.capacity() {
139            self.reserve(1)?;
140        }
141        self.output.bytes.push(byte);
142        Ok(())
143    }
144    #[inline(always)]
145    fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
146        if bytes.len() > self.output.limit - self.output.bytes.len() {
147            return Err(postcard::Error::SerializeBufferFull);
148        }
149        if bytes.len() > self.output.bytes.capacity() - self.output.bytes.len() {
150            self.reserve(bytes.len())?;
151        }
152        self.output.bytes.extend_from_slice(bytes);
153        Ok(())
154    }
155    fn finalize(self) -> postcard::Result<Self::Output> {
156        Ok(self.output.bytes)
157    }
158}
159
160// Count with the same admitted byte ceiling as encoding, without output storage.
161struct Size {
162    bytes: usize,
163    limit: usize,
164}
165impl Flavor for Size {
166    type Output = usize;
167    fn try_push(&mut self, _: u8) -> postcard::Result<()> {
168        self.try_extend(&[0])
169    }
170    fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
171        if bytes.len() > self.limit - self.bytes {
172            return Err(postcard::Error::SerializeBufferFull);
173        }
174        self.bytes += bytes.len();
175        Ok(())
176    }
177    fn finalize(self) -> postcard::Result<usize> {
178        Ok(self.bytes)
179    }
180}
181
182fn encoded_len<T: Wire>(limits: Limits, value: &T) -> Result<usize, Error> {
183    let limit = admit::<T>(limits)?;
184    postcard::serialize_with_flavor(
185        value,
186        Size {
187            bytes: HEADER,
188            limit,
189        },
190    )
191    .map_err(|_| Error::Malformed)
192}
193
194// Compare canonical serialization directly with the admitted payload. Continue
195// after a mismatch so serializer errors still retain their Malformed mapping.
196pub(crate) struct Compare<'a> {
197    pub(crate) remaining: &'a [u8],
198    pub(crate) allowance: usize,
199    pub(crate) equal: bool,
200}
201impl Flavor for Compare<'_> {
202    type Output = bool;
203    #[inline(always)]
204    fn try_push(&mut self, byte: u8) -> postcard::Result<()> {
205        if self.allowance == 0 {
206            return Err(postcard::Error::SerializeBufferFull);
207        }
208        self.allowance -= 1;
209        if let Some((first, rest)) = self.remaining.split_first() {
210            self.equal &= *first == byte;
211            self.remaining = rest;
212        } else {
213            self.equal = false;
214        }
215        Ok(())
216    }
217    fn try_extend(&mut self, bytes: &[u8]) -> postcard::Result<()> {
218        if let [byte] = bytes {
219            return self.try_push(*byte);
220        }
221        if bytes.len() > self.allowance {
222            return Err(postcard::Error::SerializeBufferFull);
223        }
224        self.allowance -= bytes.len();
225        if let Some((prefix, rest)) = self.remaining.split_at_checked(bytes.len()) {
226            // Wasm's bytewise memcmp is expensive for bulk byte/text payloads.
227            // Exact unaligned word loads preserve canonical equality; no hash or
228            // scratch allocation is involved. Small scalar writes stay simple.
229            self.equal &= if cfg!(any(target_arch = "wasm32", test))
230                && core::ptr::eq(prefix.as_ptr(), bytes.as_ptr())
231            {
232                // Borrowed decoding can serialize the exact input region.
233                // Its length and allowance were checked above; equality needs no scan.
234                true
235            } else if cfg!(any(target_arch = "wasm32", test)) && bytes.len() >= 8 {
236                if bytes.len() >= 64 {
237                    let (left, left_tail) = prefix.as_chunks::<16>();
238                    let (right, right_tail) = bytes.as_chunks::<16>();
239                    left.iter().zip(right).fold(0_u128, |difference, (a, b)| {
240                        difference | (u128::from_ne_bytes(*a) ^ u128::from_ne_bytes(*b))
241                    }) == 0
242                        && left_tail == right_tail
243                } else {
244                    let (left, left_tail) = prefix.as_chunks::<8>();
245                    let (right, right_tail) = bytes.as_chunks::<8>();
246                    left.iter().zip(right).fold(0_u64, |difference, (a, b)| {
247                        difference | (u64::from_ne_bytes(*a) ^ u64::from_ne_bytes(*b))
248                    }) == 0
249                        && left_tail == right_tail
250                }
251            } else {
252                prefix == bytes
253            };
254            self.remaining = rest;
255        } else {
256            self.equal = false;
257            self.remaining = &[];
258        }
259        Ok(())
260    }
261    fn finalize(self) -> postcard::Result<Self::Output> {
262        Ok(self.equal && self.remaining.is_empty())
263    }
264}
265
266fn encode<T: Wire>(id: u64, direction: u8, limits: Limits, value: &T) -> Result<Vec<u8>, Error> {
267    let size = admit::<T>(limits)?;
268    // The fixed header always fits after admission. Initialize the same bounded
269    // starting capacity directly, leaving the cold growth path for payloads.
270    let mut bytes = Vec::new();
271    bytes
272        .try_reserve_exact(size.min(64))
273        .map_err(|_| Error::Allocation)?;
274    bytes.extend_from_slice(&header(id, direction));
275    let allocation_failed = Cell::new(false);
276    let output = GrowingOutput {
277        output: Output { bytes, limit: size },
278        allocation_failed: &allocation_failed,
279    };
280    postcard::serialize_with_flavor(value, output).map_err(|_| {
281        if allocation_failed.get() {
282            Error::Allocation
283        } else {
284            Error::Malformed
285        }
286    })
287}
288
289fn decode<'de, T: Wire + serde::Deserialize<'de>>(
290    id: u64,
291    direction: u8,
292    limits: Limits,
293    bytes: &'de [u8],
294) -> Result<T, Error> {
295    let size = admit::<T>(limits)?;
296    if bytes.len() > size || bytes.len() > limits.bytes {
297        return Err(Error::Limit);
298    }
299    let payload = payload(id, direction, bytes)?;
300    let (value, rest) = postcard::take_from_bytes::<T>(payload).map_err(|_| Error::Malformed)?;
301    if !rest.is_empty() {
302        return Err(Error::Trailing);
303    }
304    let equal = postcard::serialize_with_flavor(
305        &value,
306        Compare {
307            remaining: payload,
308            allowance: size - HEADER,
309            equal: true,
310        },
311    )
312    .map_err(|_| Error::Malformed)?;
313    if !equal {
314        return Err(Error::Noncanonical);
315    }
316    Ok(value)
317}
318
319/// Encode once, retaining exact request bytes bound to `M`.
320pub fn encode_request<M: Method>(value: &M::Request) -> Result<RequestBytes<M>, Error> {
321    Ok(RequestBytes {
322        bytes: encode(M::ID, 0, M::REQUEST_LIMITS, value)?,
323        method: PhantomData,
324    })
325}
326/// Measure the complete request envelope without allocating an output buffer.
327/// Uses the same schema admission as encoding. Custom serializers must be
328/// deterministic; sizing does not retain bytes for retries.
329pub fn encoded_request_len<M: Method>(value: &M::Request) -> Result<usize, Error> {
330    encoded_len(M::REQUEST_LIMITS, value)
331}
332/// Measure the complete response envelope without allocating an output buffer.
333/// Consumers can compare this with a delivery allowance before encoding. Schema
334/// admission still applies to the entire type, not just this value.
335pub fn encoded_response_len<M: Method>(value: &M::Response) -> Result<usize, Error> {
336    encoded_len(M::RESPONSE_LIMITS, value)
337}
338/// Decode a request only after method, direction, canonical and resource checks.
339pub fn decode_request<'de, M: Method>(bytes: &'de [u8]) -> Result<M::Request, Error>
340where
341    M::Request: serde::Deserialize<'de>,
342{
343    decode(M::ID, 0, M::REQUEST_LIMITS, bytes)
344}
345/// Encode a reply. Application errors belong inside `M::Response`.
346pub fn encode_response<M: Method>(value: &M::Response) -> Result<ResponseBytes<M>, Error> {
347    Ok(ResponseBytes {
348        bytes: encode(M::ID, 1, M::RESPONSE_LIMITS, value)?,
349        method: PhantomData,
350    })
351}
352/// Decode a reply under the selected method's response contract.
353pub fn decode_response<'de, M: Method>(bytes: &'de [u8]) -> Result<M::Response, Error>
354where
355    M::Response: serde::Deserialize<'de>,
356{
357    decode(M::ID, 1, M::RESPONSE_LIMITS, bytes)
358}
359
360pub(crate) fn header(id: u64, direction: u8) -> [u8; HEADER] {
361    let mut header = [0; HEADER];
362    header[..2].copy_from_slice(b"IW");
363    header[2..10].copy_from_slice(&id.to_le_bytes());
364    header[10] = direction;
365    header
366}
367pub(crate) fn payload(id: u64, direction: u8, bytes: &[u8]) -> Result<&[u8], Error> {
368    if bytes.len() < HEADER {
369        return Err(Error::Malformed);
370    }
371    if &bytes[..2] != b"IW" || bytes[2..10] != id.to_le_bytes() || bytes[10] != direction {
372        return Err(Error::Identity);
373    }
374    Ok(&bytes[HEADER..])
375}
ic_wire/contract/mod.rs

ic_wire/contract/
mod.rs
