1use std::collections::BTreeSet;
25
26use crate::budget::Budget;
27use crate::cap::{CapAttrs, CapId};
28use crate::error::{Error, Result};
29use crate::promise::PromiseId;
30use crate::value::{Arg, Value};
31
32pub const MAGIC: [u8; 4] = *b"ARP1";
34pub const VERSION: u8 = 1;
36pub const FLAG_REDACT: u16 = 1 << 0;
38pub const FLAG_READ_ONLY: u16 = 1 << 1;
40pub const FLAG_KNOWN: u16 = FLAG_REDACT | FLAG_READ_ONLY;
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum Kind {
46 Request,
48 Response,
50 Error,
52 Hello,
54}
55
56impl Kind {
57 pub const fn code(self) -> u8 {
59 match self {
60 Kind::Request => 1,
61 Kind::Response => 2,
62 Kind::Error => 3,
63 Kind::Hello => 4,
64 }
65 }
66
67 pub const fn from_code(code: u8) -> Option<Kind> {
69 match code {
70 1 => Some(Kind::Request),
71 2 => Some(Kind::Response),
72 3 => Some(Kind::Error),
73 4 => Some(Kind::Hello),
74 _ => None,
75 }
76 }
77}
78
79#[derive(Debug, Clone, PartialEq)]
81pub struct CapDescriptor {
82 pub id: CapId,
84 pub methods: BTreeSet<String>,
86 pub budget: Budget,
88 pub read_only: bool,
90 pub redact: bool,
92 pub path_scope: Option<String>,
94}
95
96impl CapDescriptor {
97 pub fn from_attrs(id: CapId, attrs: &CapAttrs) -> Self {
99 Self {
100 id,
101 methods: attrs.methods.clone(),
102 budget: attrs.budget,
103 read_only: attrs.read_only,
104 redact: attrs.redact,
105 path_scope: attrs.path_scope.clone(),
106 }
107 }
108}
109
110#[derive(Debug, Clone, PartialEq)]
112pub struct Envelope {
113 pub kind: Kind,
115 pub request_id: u64,
117 pub method: String,
119 pub args: Vec<Arg>,
121 pub payload: Vec<u8>,
123 pub caps: Vec<CapDescriptor>,
125 pub redact: bool,
127 pub read_only: bool,
129}
130
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub struct WireLimits {
134 pub max_bytes: usize,
136 pub max_args: u32,
138 pub max_method: u32,
140 pub max_payload: u32,
142 pub max_caps: u32,
144 pub max_methods: u32,
146}
147
148impl Default for WireLimits {
149 fn default() -> Self {
150 Self {
151 max_bytes: 64 << 20,
152 max_args: 1 << 12,
153 max_method: 256,
154 max_payload: 32 << 20,
155 max_caps: 1 << 8,
156 max_methods: 64,
157 }
158 }
159}
160
161impl WireLimits {
162 pub fn strict() -> Self {
164 Self {
165 max_bytes: 1 << 20,
166 max_args: 64,
167 max_method: 64,
168 max_payload: 1 << 20,
169 max_caps: 16,
170 max_methods: 16,
171 }
172 }
173}
174
175impl Envelope {
176 pub fn request(request_id: u64, method: impl Into<String>, args: Vec<Arg>) -> Self {
178 Self {
179 kind: Kind::Request,
180 request_id,
181 method: method.into(),
182 args,
183 payload: Vec::new(),
184 caps: Vec::new(),
185 redact: false,
186 read_only: false,
187 }
188 }
189
190 pub fn response(request_id: u64, value: Value) -> Self {
192 Self {
193 kind: Kind::Response,
194 request_id,
195 method: String::new(),
196 args: vec![Arg::Value(value)],
197 payload: Vec::new(),
198 caps: Vec::new(),
199 redact: false,
200 read_only: false,
201 }
202 }
203
204 pub fn error(request_id: u64, error: &Error) -> Self {
206 Self {
207 kind: Kind::Error,
208 request_id,
209 method: error.code().to_string(),
210 args: Vec::new(),
211 payload: Vec::new(),
212 caps: Vec::new(),
213 redact: false,
214 read_only: false,
215 }
216 }
217
218 pub fn with_payload(mut self, payload: Vec<u8>) -> Self {
220 self.payload = payload;
221 self
222 }
223
224 pub fn with_caps(mut self, caps: Vec<CapDescriptor>) -> Self {
226 self.caps = caps;
227 self
228 }
229
230 pub fn encode(&self) -> Result<Vec<u8>> {
232 let mut out = Vec::with_capacity(32 + self.payload.len());
233 out.extend_from_slice(&MAGIC);
234 out.push(VERSION);
235 out.push(self.kind.code());
236 let mut flags = 0u16;
237 if self.redact {
238 flags |= FLAG_REDACT;
239 }
240 if self.read_only {
241 flags |= FLAG_READ_ONLY;
242 }
243 out.extend_from_slice(&flags.to_le_bytes());
244 out.extend_from_slice(&self.request_id.to_le_bytes());
245 write_str(&mut out, &self.method)?;
246 if self.args.len() > u32::MAX as usize {
247 return Err(Error::Wire("too many arguments"));
248 }
249 out.extend_from_slice(&(self.args.len() as u32).to_le_bytes());
250 for arg in &self.args {
251 write_arg(&mut out, arg)?;
252 }
253 if self.payload.len() > u32::MAX as usize {
254 return Err(Error::Wire("payload too large"));
255 }
256 out.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
257 out.extend_from_slice(&self.payload);
258 out.extend_from_slice(&(self.caps.len() as u32).to_le_bytes());
259 for cap in &self.caps {
260 write_cap(&mut out, cap)?;
261 }
262 Ok(out)
263 }
264
265 pub fn decode(bytes: &[u8]) -> Result<Envelope> {
267 Self::decode_with(bytes, WireLimits::default())
268 }
269
270 pub fn decode_with(bytes: &[u8], limits: WireLimits) -> Result<Envelope> {
272 if bytes.len() > limits.max_bytes {
273 return Err(Error::Wire("envelope exceeds the size ceiling"));
274 }
275 if bytes.len() < 16 || bytes[0..4] != MAGIC {
276 return Err(Error::Wire("bad magic"));
277 }
278 if bytes[4] != VERSION {
279 return Err(Error::Wire("unsupported envelope version"));
280 }
281 let kind = Kind::from_code(bytes[5]).ok_or(Error::Wire("unknown envelope kind"))?;
282 let flags = u16::from_le_bytes([bytes[6], bytes[7]]);
283 if flags & !FLAG_KNOWN != 0 {
284 return Err(Error::Wire("unknown envelope flags"));
285 }
286 let request_id = u64::from_le_bytes(bytes[8..16].try_into().expect("8 bytes"));
287 let mut cursor = 16usize;
288 let method = read_str(bytes, &mut cursor, limits.max_method)?;
289 let arg_count = read_u32(bytes, &mut cursor)?;
290 if arg_count > limits.max_args {
291 return Err(Error::Wire("argument count exceeds the ceiling"));
292 }
293 let mut args = Vec::with_capacity(arg_count as usize);
294 for _ in 0..arg_count {
295 args.push(read_arg(bytes, &mut cursor)?);
296 }
297 let payload_len = read_u32(bytes, &mut cursor)?;
298 if payload_len > limits.max_payload {
299 return Err(Error::Wire("payload exceeds the ceiling"));
300 }
301 let payload_end = cursor
302 .checked_add(payload_len as usize)
303 .ok_or(Error::Wire("payload length overflow"))?;
304 if payload_end > bytes.len() {
305 return Err(Error::Wire("truncated payload"));
306 }
307 let payload = bytes[cursor..payload_end].to_vec();
308 cursor = payload_end;
309 let cap_count = read_u32(bytes, &mut cursor)?;
310 if cap_count > limits.max_caps {
311 return Err(Error::CapabilityTable("too many descriptors"));
312 }
313 let mut caps = Vec::with_capacity(cap_count as usize);
314 for _ in 0..cap_count {
315 caps.push(read_cap(bytes, &mut cursor, limits.max_methods)?);
316 }
317 if cursor != bytes.len() {
318 return Err(Error::Wire("trailing bytes"));
319 }
320 Ok(Envelope {
321 kind,
322 request_id,
323 method,
324 args,
325 payload,
326 caps,
327 redact: flags & FLAG_REDACT != 0,
328 read_only: flags & FLAG_READ_ONLY != 0,
329 })
330 }
331
332 pub fn first_value(&self) -> Option<&Value> {
334 self.args.iter().find_map(|arg| match arg {
335 Arg::Value(value) => Some(value),
336 _ => None,
337 })
338 }
339}
340
341fn write_str(out: &mut Vec<u8>, text: &str) -> Result<()> {
342 if text.len() > u32::MAX as usize {
343 return Err(Error::Wire("string too large"));
344 }
345 out.extend_from_slice(&(text.len() as u32).to_le_bytes());
346 out.extend_from_slice(text.as_bytes());
347 Ok(())
348}
349
350fn write_opt_str(out: &mut Vec<u8>, text: &Option<String>) -> Result<()> {
351 match text {
352 None => {
353 out.push(0);
354 Ok(())
355 }
356 Some(text) => {
357 out.push(1);
358 write_str(out, text)
359 }
360 }
361}
362
363fn write_arg(out: &mut Vec<u8>, arg: &Arg) -> Result<()> {
364 match arg {
365 Arg::Value(value) => {
366 out.push(0);
367 write_value(out, value)
368 }
369 Arg::Promise(id) => {
370 out.push(1);
371 out.extend_from_slice(&id.0.to_le_bytes());
372 Ok(())
373 }
374 Arg::Cap(id) => {
375 out.push(2);
376 out.extend_from_slice(&id.0);
377 Ok(())
378 }
379 }
380}
381
382fn write_value(out: &mut Vec<u8>, value: &Value) -> Result<()> {
383 match value {
384 Value::Null => out.push(0),
385 Value::Bool(flag) => {
386 out.push(1);
387 out.push(u8::from(*flag));
388 }
389 Value::Int(number) => {
390 out.push(2);
391 out.extend_from_slice(&number.to_le_bytes());
392 }
393 Value::Float(number) => {
394 out.push(3);
395 out.extend_from_slice(&number.to_bits().to_le_bytes());
396 }
397 Value::Text(text) => {
398 out.push(4);
399 write_str(out, text)?;
400 }
401 Value::List(items) => {
402 out.push(5);
403 out.extend_from_slice(&(items.len() as u32).to_le_bytes());
404 for item in items {
405 write_value(out, item)?;
406 }
407 }
408 Value::Map(entries) => {
409 out.push(6);
410 out.extend_from_slice(&(entries.len() as u32).to_le_bytes());
411 for (key, value) in entries {
412 write_str(out, key)?;
413 write_value(out, value)?;
414 }
415 }
416 Value::Secret => out.push(7),
417 Value::Module(bytes) => {
418 out.push(8);
419 out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
420 out.extend_from_slice(bytes);
421 }
422 Value::Path(path) => {
423 out.push(9);
424 write_str(out, path)?;
425 }
426 }
427 Ok(())
428}
429
430fn write_cap(out: &mut Vec<u8>, cap: &CapDescriptor) -> Result<()> {
431 out.extend_from_slice(&cap.id.0);
432 out.extend_from_slice(&(cap.methods.len() as u32).to_le_bytes());
433 for method in &cap.methods {
434 write_str(out, method)?;
435 }
436 write_budget(out, &cap.budget);
437 out.push(u8::from(cap.read_only));
438 out.push(u8::from(cap.redact));
439 write_opt_str(out, &cap.path_scope)
440}
441
442fn write_budget(out: &mut Vec<u8>, budget: &Budget) {
443 for value in [
444 budget.max_calls,
445 budget.max_bytes,
446 budget.max_wall_ms,
447 budget.max_nodes,
448 budget.max_depth,
449 budget.max_array_len,
450 budget.max_value_size,
451 ] {
452 out.extend_from_slice(&value.to_le_bytes());
453 }
454}
455
456fn read_u8(bytes: &[u8], cursor: &mut usize) -> Result<u8> {
457 if *cursor >= bytes.len() {
458 return Err(Error::Wire("truncated"));
459 }
460 let value = bytes[*cursor];
461 *cursor += 1;
462 Ok(value)
463}
464
465fn read_u32(bytes: &[u8], cursor: &mut usize) -> Result<u32> {
466 if *cursor + 4 > bytes.len() {
467 return Err(Error::Wire("truncated"));
468 }
469 let value = u32::from_le_bytes(bytes[*cursor..*cursor + 4].try_into().expect("4 bytes"));
470 *cursor += 4;
471 Ok(value)
472}
473
474fn read_u64(bytes: &[u8], cursor: &mut usize) -> Result<u64> {
475 if *cursor + 8 > bytes.len() {
476 return Err(Error::Wire("truncated"));
477 }
478 let value = u64::from_le_bytes(bytes[*cursor..*cursor + 8].try_into().expect("8 bytes"));
479 *cursor += 8;
480 Ok(value)
481}
482
483fn read_str(bytes: &[u8], cursor: &mut usize, max: u32) -> Result<String> {
484 let len = read_u32(bytes, cursor)?;
485 if len > max {
486 return Err(Error::Wire("string exceeds the ceiling"));
487 }
488 let end = cursor
489 .checked_add(len as usize)
490 .ok_or(Error::Wire("length overflow"))?;
491 if end > bytes.len() {
492 return Err(Error::Wire("truncated string"));
493 }
494 let text = std::str::from_utf8(&bytes[*cursor..end])
495 .map_err(|_| Error::Wire("string is not UTF-8"))?
496 .to_string();
497 *cursor = end;
498 Ok(text)
499}
500
501fn read_opt_str(bytes: &[u8], cursor: &mut usize, max: u32) -> Result<Option<String>> {
502 match read_u8(bytes, cursor)? {
503 0 => Ok(None),
504 1 => Ok(Some(read_str(bytes, cursor, max)?)),
505 _ => Err(Error::CapabilityTable("invalid optional string tag")),
506 }
507}
508
509fn read_arg(bytes: &[u8], cursor: &mut usize) -> Result<Arg> {
510 let tag = read_u8(bytes, cursor)?;
511 match tag {
512 0 => Ok(Arg::Value(read_value(bytes, cursor)?)),
513 1 => {
514 let id = read_u64(bytes, cursor)?;
515 Ok(Arg::Promise(PromiseId(id)))
516 }
517 2 => {
518 let mut id = [0u8; 16];
519 if *cursor + 16 > bytes.len() {
520 return Err(Error::Wire("truncated capability id"));
521 }
522 id.copy_from_slice(&bytes[*cursor..*cursor + 16]);
523 *cursor += 16;
524 Ok(Arg::Cap(CapId(id)))
525 }
526 _ => Err(Error::Wire("unknown argument tag")),
527 }
528}
529
530pub const MAX_VALUE_DEPTH: usize = 64;
532
533fn read_value(bytes: &[u8], cursor: &mut usize) -> Result<Value> {
534 read_value_at(bytes, cursor, 0, bytes.len())
535}
536
537fn read_value_at(bytes: &[u8], cursor: &mut usize, depth: usize, budget: usize) -> Result<Value> {
538 if depth > MAX_VALUE_DEPTH {
539 return Err(Error::Wire("value nesting exceeds the ceiling"));
540 }
541 let tag = read_u8(bytes, cursor)?;
542 let value = match tag {
543 0 => Value::Null,
544 1 => Value::Bool(read_u8(bytes, cursor)? != 0),
545 2 => Value::Int(read_u64(bytes, cursor)? as i64),
546 3 => Value::Float(f64::from_bits(read_u64(bytes, cursor)?)),
547 4 => Value::Text(read_str(bytes, cursor, budget as u32)?),
548 5 => {
549 let count = read_u32(bytes, cursor)?;
550 if count as usize > budget {
552 return Err(Error::Wire("list too long"));
553 }
554 let mut items = Vec::with_capacity((count as usize).min(1024));
557 for _ in 0..count {
558 items.push(read_value_at(bytes, cursor, depth + 1, budget)?);
559 }
560 Value::List(items)
561 }
562 6 => {
563 let count = read_u32(bytes, cursor)?;
564 if count as usize > budget {
565 return Err(Error::Wire("map too large"));
566 }
567 let mut entries = std::collections::BTreeMap::new();
568 for _ in 0..count {
569 let key = read_str(bytes, cursor, budget as u32)?;
570 entries.insert(key, read_value_at(bytes, cursor, depth + 1, budget)?);
571 }
572 Value::Map(entries)
573 }
574 7 => Value::Secret,
575 8 => {
576 let len = read_u32(bytes, cursor)?;
577 if len as usize > budget {
578 return Err(Error::Wire("module payload exceeds the ceiling"));
579 }
580 let end = cursor
581 .checked_add(len as usize)
582 .ok_or(Error::Wire("length overflow"))?;
583 if end > bytes.len() {
584 return Err(Error::Wire("truncated module"));
585 }
586 let module = bytes[*cursor..end].to_vec();
587 *cursor = end;
588 Value::Module(module)
589 }
590 9 => Value::Path(read_str(bytes, cursor, budget as u32)?),
591 _ => return Err(Error::Wire("unknown value tag")),
592 };
593 Ok(value)
594}
595
596fn read_budget(bytes: &[u8], cursor: &mut usize) -> Result<Budget> {
597 Ok(Budget {
598 max_calls: read_u64(bytes, cursor)?,
599 max_bytes: read_u64(bytes, cursor)?,
600 max_wall_ms: read_u64(bytes, cursor)?,
601 max_nodes: read_u64(bytes, cursor)?,
602 max_depth: read_u64(bytes, cursor)?,
603 max_array_len: read_u64(bytes, cursor)?,
604 max_value_size: read_u64(bytes, cursor)?,
605 })
606}
607
608fn read_cap(bytes: &[u8], cursor: &mut usize, max_methods: u32) -> Result<CapDescriptor> {
609 let mut id = [0u8; 16];
610 if *cursor + 16 > bytes.len() {
611 return Err(Error::CapabilityTable("truncated capability id"));
612 }
613 id.copy_from_slice(&bytes[*cursor..*cursor + 16]);
614 *cursor += 16;
615 let count = read_u32(bytes, cursor)?;
616 if count > max_methods {
617 return Err(Error::CapabilityTable("too many methods"));
618 }
619 let mut methods = BTreeSet::new();
620 for _ in 0..count {
621 methods.insert(read_str(bytes, cursor, 256)?);
622 }
623 let budget = read_budget(bytes, cursor)?;
624 budget
625 .check()
626 .map_err(|_| Error::CapabilityTable("budget disables a ceiling"))?;
627 let read_only = read_u8(bytes, cursor)? != 0;
628 let redact = read_u8(bytes, cursor)? != 0;
629 let path_scope = read_opt_str(bytes, cursor, 1024)?;
630 Ok(CapDescriptor {
631 id: CapId(id),
632 methods,
633 budget,
634 read_only,
635 redact,
636 path_scope,
637 })
638}
639
640pub fn decode_capability_table(bytes: &[u8]) -> Result<Vec<CapDescriptor>> {
642 let mut cursor = 0usize;
643 let count = read_u32(bytes, &mut cursor)?;
644 if count > WireLimits::default().max_caps {
645 return Err(Error::CapabilityTable("too many descriptors"));
646 }
647 let mut caps = Vec::with_capacity((count as usize).min(64));
648 for _ in 0..count {
649 caps.push(read_cap(bytes, &mut cursor, 256)?);
650 }
651 if cursor != bytes.len() {
652 return Err(Error::CapabilityTable("trailing bytes"));
653 }
654 Ok(caps)
655}
656
657pub fn encode_capability_table(caps: &[CapDescriptor]) -> Result<Vec<u8>> {
659 let mut out = Vec::new();
660 out.extend_from_slice(&(caps.len() as u32).to_le_bytes());
661 for cap in caps {
662 write_cap(&mut out, cap)?;
663 }
664 Ok(out)
665}