libcfd_rpc/
io.rs

1use core::future::poll_fn;
2use core::pin::Pin;
3
4use futures_io::{AsyncRead, AsyncWrite};
5
6use crate::error::{Result, RpcError};
7
8/// A bidirectional byte stream capable of carrying RPC messages.
9///
10/// This mirrors the transport cloudflared uses for its registration stream:
11/// a single full-duplex stream (QUIC stream or HTTP/2 request/response body).
12pub trait AsyncStream: AsyncRead + AsyncWrite + Send + Unpin {}
13
14impl<T> AsyncStream for T where T: AsyncRead + AsyncWrite + Send + Unpin {}
15
16/// Reads one Cap'n Proto message (including its segment table) from the stream.
17///
18/// Framing: `[u32 LE (numSegments-1)] [per segment: u32 LE word count]
19/// [zero u32 pad to 8-byte alignment] [segment data]`.
20///
21/// Mirrors capnp-go's stream decoder including its `defaultDecodeLimit`
22/// of 64 MiB of total segment data.
23pub async fn read_message<S: AsyncStream + Unpin>(
24    stream: &mut S,
25) -> Result<capnp::message::Reader<capnp::serialize::OwnedSegments>> {
26    let mut header = [0u8; 8];
27    read_exact(stream, &mut header).await?;
28    let segment_count =
29        u32::from_le_bytes(header[0..4].try_into().expect("4 bytes")).wrapping_add(1) as usize;
30    if segment_count == 0 || segment_count > 512 {
31        return Err(RpcError::Protocol(format!(
32            "invalid segment count {segment_count}"
33        )));
34    }
35
36    let mut builder = capnp::serialize::SegmentLengthsBuilder::with_capacity(segment_count);
37    let mut total_words = 0usize;
38    let first = u32::from_le_bytes(header[4..8].try_into().expect("4 bytes")) as usize;
39    total_words = total_words.saturating_add(first);
40    if total_words > MAXIMUM_TOTAL_WORDS {
41        return Err(RpcError::Protocol(format!(
42            "message too large: {total_words} words exceeds limit of {MAXIMUM_TOTAL_WORDS}"
43        )));
44    }
45    builder.try_push_segment(first)?;
46    if segment_count > 1 {
47        // Go: streamHeaderSize(maxSeg) = (4 + 4*numSegments + 7) & !7
48        let header_size = (4 + 4 * segment_count + 7) & !7;
49        let mut sizes = vec![0u8; header_size - 8];
50        read_exact(stream, &mut sizes).await?;
51        for i in 1..segment_count {
52            let offset = (i - 1) * 4;
53            let length =
54                u32::from_le_bytes(sizes[offset..offset + 4].try_into().expect("4 bytes")) as usize;
55            total_words = total_words.saturating_add(length);
56            if total_words > MAXIMUM_TOTAL_WORDS {
57                return Err(RpcError::Protocol(format!(
58                    "message too large: {total_words} words exceeds limit of {MAXIMUM_TOTAL_WORDS}"
59                )));
60            }
61            builder.try_push_segment(length)?;
62        }
63    }
64
65    let mut segments = builder.into_owned_segments();
66    read_exact(stream, &mut segments[..]).await?;
67    Ok(capnp::message::Reader::new(
68        segments,
69        capnp::message::ReaderOptions::new(),
70    ))
71}
72
73/// capnp-go's `defaultDecodeLimit` is 64 MiB; as words (8 bytes each) with
74/// one header word per segment this bounds a single message.
75const MAXIMUM_TOTAL_WORDS: usize = 8 * 1024 * 1024;
76
77/// Serializes a Cap'n Proto message (including its segment table) to framed
78/// bytes. Synchronous so no non-`Send` capnp builder state is held across an
79/// await; write the result with [`write_raw`].
80pub fn serialize_message(
81    message: &capnp::message::Builder<capnp::message::HeapAllocator>,
82) -> Vec<u8> {
83    capnp::serialize::write_message_to_words(message)
84}
85
86/// Writes already-framed bytes (as produced by `serialize::write_message_to_words`)
87/// to the stream.
88pub async fn write_raw<S: AsyncStream + Unpin>(stream: &mut S, bytes: &[u8]) -> Result<()> {
89    write_all(stream, bytes).await?;
90    Ok(())
91}
92
93async fn read_exact<S: AsyncRead + Unpin>(stream: &mut S, mut buffer: &mut [u8]) -> Result<()> {
94    while !buffer.is_empty() {
95        let n = poll_fn(|cx| Pin::new(&mut *stream).poll_read(cx, buffer)).await?;
96        if n == 0 {
97            return Err(RpcError::Eof);
98        }
99        let tmp = buffer;
100        buffer = &mut tmp[n..];
101    }
102    Ok(())
103}
104
105async fn write_all<S: AsyncWrite + Unpin>(stream: &mut S, mut buffer: &[u8]) -> Result<()> {
106    while !buffer.is_empty() {
107        let n = poll_fn(|cx| Pin::new(&mut *stream).poll_write(cx, buffer)).await?;
108        if n == 0 {
109            return Err(RpcError::Protocol("write returned zero bytes".into()));
110        }
111        buffer = &buffer[n..];
112    }
113    Ok(())
114}
libcfd_rpc/io.rs

libcfd_rpc/
io.rs
