Skip to main content

nub_program/
codec.rs

1//! Binary encoding for [`ProgramBlob`].
2//!
3//! A deliberately dumb little-endian format with no dependencies — not
4//! SSZ, not serde. This artifact is a build product consumed by a
5//! runtime in the same tree, so it needs neither content-addressing nor
6//! a stable cross-ecosystem schema; a personality that wants those
7//! wraps the blob in its own encoding.
8//!
9//! ```text
10//!   magic     "NUBP"                                     4
11//!   version   u16                                        2
12//!   flags     u16 (reserved, must be 0)                   2
13//!   stack_pages ro_pages rw_pages heap_pages   u32 x4    16
14//!   code_len ro_len rw_len endpoint_count      u32 x4    16
15//!   endpoints[endpoint_count]:
16//!       key u8 | arg_registers u8 | arg_meta u8 | reg_count u8
17//!       entry_pc u64
18//!       reg_count x (idx u8 | pad u8 x7 | value u64)
19//!   code[code_len] ro[ro_len] rw[rw_len]
20//! ```
21//!
22//! `ro_len`/`rw_len` are the *trailing-zero-trimmed* lengths; decode
23//! zero-extends each back to `pages * PAGE_SIZE`. That keeps `.bss`-
24//! heavy programs (the 64 KiB guest bump arenas) from paying for their
25//! zeros on disk, and is why [`ProgramBlob::new`] normalizes the
26//! buffers to whole pages: trim-then-extend then round-trips exactly.
27
28use alloc::collections::BTreeMap;
29use alloc::vec::Vec;
30
31use crate::abi::PAGE_SIZE;
32use crate::blob::{Endpoint, InvalidProgram, ProgramBlob, Regions};
33
34/// Format magic: `b"NUBP"`.
35pub const MAGIC: [u8; 4] = *b"NUBP";
36/// Current format version.
37pub const VERSION: u16 = 1;
38
39const HEADER_LEN: usize = 4 + 2 + 2 + 16 + 16;
40const ENDPOINT_HEAD_LEN: usize = 4 + 8;
41const REG_ENTRY_LEN: usize = 8 + 8;
42
43/// Why a byte slice does not decode to a [`ProgramBlob`].
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub enum DecodeError {
46    /// The leading 4 bytes are not [`MAGIC`].
47    BadMagic,
48    /// Encoded by a newer (or older, incompatible) writer.
49    UnsupportedVersion(u16),
50    /// A reserved header field was non-zero.
51    ReservedFlags(u16),
52    /// The input ended mid-field.
53    Truncated { need: usize, have: usize },
54    /// Two endpoint records claim the same index.
55    DuplicateEndpoint(u8),
56    /// A trimmed region length exceeds its page count.
57    RegionOverflow { len: u32, capacity: usize },
58    /// Trailing bytes after the last declared field.
59    TrailingBytes(usize),
60    /// Decoded successfully but violates a [`ProgramBlob`] invariant.
61    Invalid(InvalidProgram),
62}
63
64impl core::fmt::Display for DecodeError {
65    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66        match self {
67            DecodeError::BadMagic => f.write_str("not a nub program blob (bad magic)"),
68            DecodeError::UnsupportedVersion(v) => {
69                write!(
70                    f,
71                    "unsupported blob version {v} (this build reads {VERSION})"
72                )
73            }
74            DecodeError::ReservedFlags(v) => {
75                write!(f, "reserved flags field is {v:#x}, expected 0")
76            }
77            DecodeError::Truncated { need, have } => {
78                write!(f, "truncated: need {need} bytes, have {have}")
79            }
80            DecodeError::DuplicateEndpoint(k) => write!(f, "duplicate endpoint index {k}"),
81            DecodeError::RegionOverflow { len, capacity } => write!(
82                f,
83                "region payload {len} bytes exceeds its {capacity}-byte page capacity"
84            ),
85            DecodeError::TrailingBytes(n) => write!(f, "{n} trailing bytes after the blob"),
86            DecodeError::Invalid(e) => write!(f, "{e}"),
87        }
88    }
89}
90
91impl core::error::Error for DecodeError {}
92
93impl From<InvalidProgram> for DecodeError {
94    fn from(e: InvalidProgram) -> Self {
95        DecodeError::Invalid(e)
96    }
97}
98
99/// Length of `data` with trailing zero bytes removed.
100fn trimmed_len(data: &[u8]) -> usize {
101    match data.iter().rposition(|&b| b != 0) {
102        Some(i) => i + 1,
103        None => 0,
104    }
105}
106
107/// Cursor over the input that reports how far it got when it runs out.
108struct Reader<'a> {
109    buf: &'a [u8],
110    pos: usize,
111}
112
113impl<'a> Reader<'a> {
114    fn take(&mut self, n: usize) -> Result<&'a [u8], DecodeError> {
115        let end = self.pos.checked_add(n).ok_or(DecodeError::Truncated {
116            need: usize::MAX,
117            have: self.buf.len(),
118        })?;
119        if end > self.buf.len() {
120            return Err(DecodeError::Truncated {
121                need: end,
122                have: self.buf.len(),
123            });
124        }
125        let out = &self.buf[self.pos..end];
126        self.pos = end;
127        Ok(out)
128    }
129
130    fn u8(&mut self) -> Result<u8, DecodeError> {
131        Ok(self.take(1)?[0])
132    }
133
134    fn u16(&mut self) -> Result<u16, DecodeError> {
135        Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
136    }
137
138    fn u32(&mut self) -> Result<u32, DecodeError> {
139        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
140    }
141
142    fn u64(&mut self) -> Result<u64, DecodeError> {
143        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
144    }
145}
146
147impl ProgramBlob {
148    /// Serialize to the format documented on this module.
149    pub fn to_bytes(&self) -> Vec<u8> {
150        let ro_len = trimmed_len(&self.ro_data);
151        let rw_len = trimmed_len(&self.rw_data);
152
153        let endpoints_len: usize = self
154            .endpoints
155            .values()
156            .map(|e| ENDPOINT_HEAD_LEN + e.initial_regs.len() * REG_ENTRY_LEN)
157            .sum();
158        let mut out =
159            Vec::with_capacity(HEADER_LEN + endpoints_len + self.code.len() + ro_len + rw_len);
160
161        out.extend_from_slice(&MAGIC);
162        out.extend_from_slice(&VERSION.to_le_bytes());
163        out.extend_from_slice(&0u16.to_le_bytes()); // flags
164        for v in [
165            self.regions.stack_pages,
166            self.regions.ro_pages,
167            self.regions.rw_pages,
168            self.regions.heap_pages,
169        ] {
170            out.extend_from_slice(&v.to_le_bytes());
171        }
172        for v in [
173            self.code.len() as u32,
174            ro_len as u32,
175            rw_len as u32,
176            self.endpoints.len() as u32,
177        ] {
178            out.extend_from_slice(&v.to_le_bytes());
179        }
180
181        for (&key, ep) in &self.endpoints {
182            out.push(key);
183            out.push(ep.arg_registers);
184            out.push(ep.arg_meta);
185            out.push(ep.initial_regs.len() as u8);
186            out.extend_from_slice(&ep.entry_pc.to_le_bytes());
187            for (&idx, &value) in &ep.initial_regs {
188                out.push(idx);
189                out.extend_from_slice(&[0u8; 7]);
190                out.extend_from_slice(&value.to_le_bytes());
191            }
192        }
193
194        out.extend_from_slice(&self.code);
195        out.extend_from_slice(&self.ro_data[..ro_len]);
196        out.extend_from_slice(&self.rw_data[..rw_len]);
197        out
198    }
199
200    /// Parse bytes produced by [`ProgramBlob::to_bytes`].
201    ///
202    /// Rejects trailing bytes: a blob is a whole file, and silently
203    /// ignoring a suffix would hide a truncated or concatenated write.
204    pub fn from_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
205        let mut r = Reader { buf: bytes, pos: 0 };
206
207        if r.take(4)? != MAGIC {
208            return Err(DecodeError::BadMagic);
209        }
210        let version = r.u16()?;
211        if version != VERSION {
212            return Err(DecodeError::UnsupportedVersion(version));
213        }
214        let flags = r.u16()?;
215        if flags != 0 {
216            return Err(DecodeError::ReservedFlags(flags));
217        }
218
219        let regions = Regions {
220            stack_pages: r.u32()?,
221            ro_pages: r.u32()?,
222            rw_pages: r.u32()?,
223            heap_pages: r.u32()?,
224        };
225        let code_len = r.u32()? as usize;
226        let ro_len = r.u32()?;
227        let rw_len = r.u32()?;
228        let endpoint_count = r.u32()?;
229
230        let mut endpoints: BTreeMap<u8, Endpoint> = BTreeMap::new();
231        for _ in 0..endpoint_count {
232            let key = r.u8()?;
233            let arg_registers = r.u8()?;
234            let arg_meta = r.u8()?;
235            let reg_count = r.u8()?;
236            let entry_pc = r.u64()?;
237            let mut initial_regs = BTreeMap::new();
238            for _ in 0..reg_count {
239                let idx = r.u8()?;
240                let _pad = r.take(7)?;
241                initial_regs.insert(idx, r.u64()?);
242            }
243            if endpoints
244                .insert(
245                    key,
246                    Endpoint {
247                        entry_pc,
248                        arg_registers,
249                        arg_meta,
250                        initial_regs,
251                    },
252                )
253                .is_some()
254            {
255                return Err(DecodeError::DuplicateEndpoint(key));
256            }
257        }
258
259        let code = r.take(code_len)?.to_vec();
260        let ro_data = read_region(&mut r, ro_len, regions.ro_pages)?;
261        let rw_data = read_region(&mut r, rw_len, regions.rw_pages)?;
262
263        if r.pos != bytes.len() {
264            return Err(DecodeError::TrailingBytes(bytes.len() - r.pos));
265        }
266
267        let blob = ProgramBlob {
268            code,
269            regions,
270            ro_data,
271            rw_data,
272            endpoints,
273        };
274        blob.validate()?;
275        Ok(blob)
276    }
277}
278
279/// Read `len` payload bytes and zero-extend to `pages * PAGE_SIZE`.
280fn read_region(r: &mut Reader<'_>, len: u32, pages: u32) -> Result<Vec<u8>, DecodeError> {
281    let capacity = pages as usize * PAGE_SIZE as usize;
282    if len as usize > capacity {
283        return Err(DecodeError::RegionOverflow { len, capacity });
284    }
285    let mut data = r.take(len as usize)?.to_vec();
286    data.resize(capacity, 0);
287    Ok(data)
288}