1use alloc::collections::BTreeMap;
29use alloc::vec::Vec;
30
31use crate::abi::PAGE_SIZE;
32use crate::blob::{Endpoint, InvalidProgram, ProgramBlob, Regions};
33
34pub const MAGIC: [u8; 4] = *b"NUBP";
36pub 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#[derive(Debug, Clone, PartialEq, Eq)]
45pub enum DecodeError {
46 BadMagic,
48 UnsupportedVersion(u16),
50 ReservedFlags(u16),
52 Truncated { need: usize, have: usize },
54 DuplicateEndpoint(u8),
56 RegionOverflow { len: u32, capacity: usize },
58 TrailingBytes(usize),
60 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
99fn trimmed_len(data: &[u8]) -> usize {
101 match data.iter().rposition(|&b| b != 0) {
102 Some(i) => i + 1,
103 None => 0,
104 }
105}
106
107struct 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 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()); 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 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
279fn 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}