Skip to main content

chio_kernel/
transport.rs

1//! Length-prefixed canonical JSON transport.
2//!
3//! Wire format: `[4-byte big-endian length][canonical JSON bytes]`
4//!
5//! The transport is generic over `Read` and `Write` so it works with pipes,
6//! TCP, Unix domain sockets, or in-memory buffers for testing.
7
8use std::io::{BufReader, BufWriter, Read, Write};
9
10use chio_core::canonical::canonical_json_bytes;
11use chio_core::message::{AgentMessage, KernelMessage};
12
13/// Errors produced by the transport layer.
14#[derive(Debug, thiserror::Error)]
15pub enum TransportError {
16    #[error("i/o error: {0}")]
17    Io(#[from] std::io::Error),
18
19    #[error("message too large: {size} bytes (max {max})")]
20    MessageTooLarge { size: u32, max: u32 },
21
22    #[error("json deserialization error: {0}")]
23    Deserialize(#[from] serde_json::Error),
24
25    #[error("canonical json serialization error: {0}")]
26    Serialize(String),
27
28    #[error("connection closed")]
29    ConnectionClosed,
30}
31
32/// Maximum message size: 16 MiB.
33const MAX_MESSAGE_SIZE: u32 = 16 * 1024 * 1024;
34
35/// Length-prefixed canonical JSON transport.
36///
37/// Reads `AgentMessage` frames from the reader and writes `KernelMessage`
38/// frames to the writer. Each frame is a 4-byte big-endian length prefix
39/// followed by that many bytes of canonical JSON.
40pub struct ChioTransport<R: Read, W: Write> {
41    reader: BufReader<R>,
42    writer: BufWriter<W>,
43}
44
45impl<R: Read, W: Write> ChioTransport<R, W> {
46    pub fn new(reader: R, writer: W) -> Self {
47        Self {
48            reader: BufReader::new(reader),
49            writer: BufWriter::new(writer),
50        }
51    }
52
53    /// Read one `AgentMessage` from the transport.
54    ///
55    /// Blocks until a complete frame is available. Returns
56    /// `TransportError::ConnectionClosed` if the reader reaches EOF before
57    /// a complete frame is read.
58    pub fn recv(&mut self) -> Result<AgentMessage, TransportError> {
59        let bytes = read_frame(&mut self.reader)?;
60        let msg: AgentMessage = serde_json::from_slice(&bytes)?;
61        Ok(msg)
62    }
63
64    /// Send one `KernelMessage` over the transport.
65    ///
66    /// The message is serialized to canonical JSON (RFC 8785) and written
67    /// as a length-prefixed frame. The writer is flushed after each send.
68    pub fn send(&mut self, msg: &KernelMessage) -> Result<(), TransportError> {
69        let bytes =
70            canonical_json_bytes(msg).map_err(|e| TransportError::Serialize(e.to_string()))?;
71        write_frame(&mut self.writer, &bytes)?;
72        self.writer.flush()?;
73        Ok(())
74    }
75}
76
77/// Read a single length-prefixed frame from a reader.
78pub fn read_frame<R: Read>(reader: &mut R) -> Result<Vec<u8>, TransportError> {
79    let mut len_buf = [0u8; 4];
80    match reader.read_exact(&mut len_buf) {
81        Ok(()) => {}
82        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
83            return Err(TransportError::ConnectionClosed);
84        }
85        Err(e) => return Err(TransportError::Io(e)),
86    }
87
88    let len = u32::from_be_bytes(len_buf);
89    if len > MAX_MESSAGE_SIZE {
90        return Err(TransportError::MessageTooLarge {
91            size: len,
92            max: MAX_MESSAGE_SIZE,
93        });
94    }
95
96    let mut buf = vec![0u8; len as usize];
97    match reader.read_exact(&mut buf) {
98        Ok(()) => {}
99        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
100            return Err(TransportError::ConnectionClosed);
101        }
102        Err(e) => return Err(TransportError::Io(e)),
103    }
104    Ok(buf)
105}
106
107/// Write a single length-prefixed frame to a writer.
108pub fn write_frame<W: Write>(writer: &mut W, data: &[u8]) -> Result<(), TransportError> {
109    let len = u32::try_from(data.len()).map_err(|_| TransportError::MessageTooLarge {
110        size: u32::MAX,
111        max: MAX_MESSAGE_SIZE,
112    })?;
113    if len > MAX_MESSAGE_SIZE {
114        return Err(TransportError::MessageTooLarge {
115            size: len,
116            max: MAX_MESSAGE_SIZE,
117        });
118    }
119    writer.write_all(&len.to_be_bytes())?;
120    writer.write_all(data)?;
121    Ok(())
122}
123
124#[cfg(test)]
125#[allow(clippy::unwrap_used, clippy::expect_used)]
126mod tests {
127    use super::*;
128    use chio_core::capability::{
129        scope::{ChioScope, Operation, ToolGrant},
130        token::{CapabilityToken, CapabilityTokenBody},
131    };
132    use chio_core::crypto::Keypair;
133    use chio_core::receipt::{
134        body::ChioReceipt, body::ChioReceiptBody, decision::Decision, decision::ToolCallAction,
135        metadata::GuardEvidence,
136    };
137    use std::io::Cursor;
138
139    fn make_token(kp: &Keypair) -> CapabilityToken {
140        let body = CapabilityTokenBody {
141            id: "cap-transport-001".to_string(),
142            issuer: kp.public_key(),
143            subject: kp.public_key(),
144            scope: ChioScope {
145                grants: vec![ToolGrant {
146                    server_id: "srv".to_string(),
147                    tool_name: "echo".to_string(),
148                    operations: vec![Operation::Invoke],
149                    constraints: vec![],
150                    max_invocations: None,
151                    max_cost_per_invocation: None,
152                    max_total_cost: None,
153                    dpop_required: None,
154                }],
155                ..ChioScope::default()
156            },
157            issued_at: 1000,
158            expires_at: 2000,
159            delegation_chain: vec![],
160            aggregate_invocation_budget: None,
161        };
162        CapabilityToken::sign(body, kp).unwrap()
163    }
164
165    fn make_receipt(kp: &Keypair) -> ChioReceipt {
166        let body = ChioReceiptBody {
167            id: "rcpt-transport-001".to_string(),
168            timestamp: 1500,
169            capability_id: "cap-transport-001".to_string(),
170            tool_server: "srv".to_string(),
171            tool_name: "echo".to_string(),
172            action: ToolCallAction::from_parameters(serde_json::json!({"text": "hello"})).unwrap(),
173            decision: Some(Decision::Allow),
174            receipt_kind: chio_core::receipt::kinds::ReceiptKind::MediatedDecision,
175            boundary_class: chio_core::receipt::kinds::BoundaryClass::Prevent,
176            observation_outcome: None,
177            tool_origin: chio_core::receipt::kinds::ToolOrigin::CallerExecuted,
178            redaction_mode: chio_core::receipt::kinds::RedactionMode::None,
179            actor_chain: Vec::new(),
180            content_hash: chio_core::sha256_hex(br#"{"output":"world"}"#),
181            policy_hash: "deadbeef".to_string(),
182            evidence: vec![GuardEvidence {
183                guard_name: "ShellCommandGuard".to_string(),
184                verdict: true,
185                details: None,
186            }],
187            metadata: None,
188            trust_level: chio_core::receipt::kinds::TrustLevel::default(),
189            tenant_id: None,
190            kernel_key: kp.public_key(),
191            bbs_projection_version: None,
192        };
193        ChioReceipt::sign(body, kp).unwrap()
194    }
195
196    #[test]
197    fn frame_roundtrip() {
198        let data = b"hello, world";
199        let mut buf = Vec::new();
200        write_frame(&mut buf, data).unwrap();
201
202        let mut cursor = Cursor::new(buf);
203        let recovered = read_frame(&mut cursor).unwrap();
204        assert_eq!(recovered, data);
205    }
206
207    #[test]
208    fn length_prefix_encoding() {
209        let data = vec![0xAA; 256];
210        let mut buf = Vec::new();
211        write_frame(&mut buf, &data).unwrap();
212
213        // First 4 bytes should be big-endian 256.
214        assert_eq!(&buf[..4], &[0, 0, 1, 0]);
215        assert_eq!(buf.len(), 4 + 256);
216    }
217
218    #[test]
219    fn transport_agent_message_roundtrip() {
220        let kp = Keypair::generate();
221        let msg = AgentMessage::ToolCallRequest {
222            id: "req-001".to_string(),
223            capability_token: Box::new(make_token(&kp)),
224            server_id: "srv".to_string(),
225            tool: "echo".to_string(),
226            params: Box::new(serde_json::json!({"text": "hello"})),
227            governed_intent: None,
228            approval_token: None,
229            approval_tokens: Vec::new(),
230            threshold_approval_proposal: None,
231            supplemental_authorization: None,
232            execution_nonce: None,
233        };
234
235        // Serialize to a buffer (using canonical JSON, same as KernelMessage path).
236        let bytes = canonical_json_bytes(&msg).expect("canonical serialization");
237        let mut wire = Vec::new();
238        write_frame(&mut wire, &bytes).unwrap();
239
240        // Read it back.
241        let mut cursor = Cursor::new(wire);
242        let frame = read_frame(&mut cursor).unwrap();
243        let recovered: AgentMessage = serde_json::from_slice(&frame).unwrap();
244
245        let (id, server_id, tool) = match recovered {
246            AgentMessage::ToolCallRequest {
247                id,
248                server_id,
249                tool,
250                ..
251            } => Some((id, server_id, tool)),
252            _ => None,
253        }
254        .expect("wrong variant");
255        assert_eq!(id, "req-001");
256        assert_eq!(server_id, "srv");
257        assert_eq!(tool, "echo");
258    }
259
260    #[test]
261    fn transport_kernel_message_roundtrip() {
262        let kp = Keypair::generate();
263        let receipt = make_receipt(&kp);
264        let kernel_msg = KernelMessage::ToolCallResponse {
265            id: "req-001".to_string(),
266            result: chio_core::message::ToolCallResult::Ok {
267                value: serde_json::json!({"output": "world"}),
268            },
269            receipt: Box::new(receipt),
270            execution_nonce: None,
271        };
272
273        // Use a shared buffer as the "pipe".
274        let mut wire = Vec::new();
275        {
276            let bytes = canonical_json_bytes(&kernel_msg).expect("canonical serialization");
277            write_frame(&mut wire, &bytes).unwrap();
278        }
279
280        let mut cursor = Cursor::new(wire);
281        let frame = read_frame(&mut cursor).unwrap();
282        let recovered: KernelMessage = serde_json::from_slice(&frame).unwrap();
283
284        let (id, result, receipt) = match recovered {
285            KernelMessage::ToolCallResponse {
286                id,
287                result,
288                receipt,
289                ..
290            } => Some((id, result, receipt)),
291            _ => None,
292        }
293        .expect("wrong variant");
294        assert_eq!(id, "req-001");
295        assert!(matches!(
296            result,
297            chio_core::message::ToolCallResult::Ok { .. }
298        ));
299        assert!(receipt.verify_signature().unwrap());
300    }
301
302    #[test]
303    fn transport_kernel_chunk_roundtrip() {
304        let kernel_msg = KernelMessage::ToolCallChunk {
305            id: "req-stream-1".to_string(),
306            chunk_index: 1,
307            data: serde_json::json!({"delta": "world"}),
308        };
309
310        let mut wire = Vec::new();
311        {
312            let bytes = canonical_json_bytes(&kernel_msg).expect("canonical serialization");
313            write_frame(&mut wire, &bytes).unwrap();
314        }
315
316        let mut cursor = Cursor::new(wire);
317        let frame = read_frame(&mut cursor).unwrap();
318        let recovered: KernelMessage = serde_json::from_slice(&frame).unwrap();
319
320        let (id, chunk_index, data) = match recovered {
321            KernelMessage::ToolCallChunk {
322                id,
323                chunk_index,
324                data,
325            } => Some((id, chunk_index, data)),
326            _ => None,
327        }
328        .expect("wrong variant");
329        assert_eq!(id, "req-stream-1");
330        assert_eq!(chunk_index, 1);
331        assert_eq!(data["delta"], "world");
332    }
333
334    #[test]
335    fn transport_send_recv_roundtrip() {
336        let agent_msg = AgentMessage::Heartbeat;
337        let agent_bytes = canonical_json_bytes(&agent_msg).expect("canonical");
338        let mut agent_wire = Vec::new();
339        write_frame(&mut agent_wire, &agent_bytes).unwrap();
340
341        let kernel_msg = KernelMessage::Heartbeat;
342
343        // agent_wire is what the "agent" wrote, kernel_buf is where the
344        // kernel writes its response.
345        let kernel_buf: Vec<u8> = Vec::new();
346        let mut transport = ChioTransport::new(Cursor::new(agent_wire), kernel_buf);
347
348        let received = transport.recv().unwrap();
349        assert!(matches!(received, AgentMessage::Heartbeat));
350
351        transport.send(&kernel_msg).unwrap();
352    }
353
354    #[test]
355    fn connection_closed_on_empty_read() {
356        let empty: Vec<u8> = Vec::new();
357        let mut cursor = Cursor::new(empty);
358        let err = read_frame(&mut cursor).unwrap_err();
359        assert!(matches!(err, TransportError::ConnectionClosed));
360    }
361
362    #[test]
363    fn rejects_oversized_frame() {
364        // Craft a length prefix claiming 20 MiB.
365        let len: u32 = 20 * 1024 * 1024;
366        let mut buf = Vec::new();
367        buf.extend_from_slice(&len.to_be_bytes());
368        buf.extend_from_slice(&[0u8; 16]); // some trailing data
369
370        let mut cursor = Cursor::new(buf);
371        let err = read_frame(&mut cursor).unwrap_err();
372        assert!(matches!(err, TransportError::MessageTooLarge { .. }));
373    }
374
375    #[test]
376    fn multiple_frames_in_sequence() {
377        let mut wire = Vec::new();
378        write_frame(&mut wire, b"first").unwrap();
379        write_frame(&mut wire, b"second").unwrap();
380        write_frame(&mut wire, b"third").unwrap();
381
382        let mut cursor = Cursor::new(wire);
383        assert_eq!(read_frame(&mut cursor).unwrap(), b"first");
384        assert_eq!(read_frame(&mut cursor).unwrap(), b"second");
385        assert_eq!(read_frame(&mut cursor).unwrap(), b"third");
386
387        // Next read should get ConnectionClosed.
388        assert!(matches!(
389            read_frame(&mut cursor).unwrap_err(),
390            TransportError::ConnectionClosed
391        ));
392    }
393}