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