Skip to main content

urma_node/
peer_codec.rs

1//! Bounded framing for untrusted peer JSON. No record or chain verification.
2//! The request-response layer owns timeouts, including waiting for stream EOF.
3use crate::peer_protocol::{
4    INVENTORY_LIMIT, PEERS_LIMIT, PROTOCOL, REQUEST_LIMIT, RESPONSE_LIMIT, Reply, Request, Response,
5};
6use async_trait::async_trait;
7use libp2p::{
8    StreamProtocol,
9    futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
10};
11use serde::{Serialize, de::DeserializeOwned};
12use std::io::{self, Write};
13
14#[derive(Clone, Default)]
15pub struct Codec;
16
17fn protocol(protocol: &StreamProtocol) -> io::Result<()> {
18    if protocol.as_ref() != PROTOCOL {
19        return Err(io::Error::new(
20            io::ErrorKind::InvalidInput,
21            "unsupported peer protocol",
22        ));
23    }
24    Ok(())
25}
26
27/// Inventory pages and peer gossip have their own, smaller frame bounds.
28fn response_limit(response: &Response) -> usize {
29    match response.result {
30        Reply::Inventory(_) => INVENTORY_LIMIT,
31        Reply::Peers(_) => PEERS_LIMIT,
32        _ => RESPONSE_LIMIT,
33    }
34}
35
36async fn read<T: AsyncRead + Unpin + Send, V: DeserializeOwned>(
37    io: &mut T,
38    limit: usize,
39) -> io::Result<(V, usize)> {
40    let mut prefix = [0; 4];
41    io.read_exact(&mut prefix).await?;
42    let length = u32::from_be_bytes(prefix) as usize;
43    if length == 0 || length > limit {
44        return Err(io::Error::new(
45            io::ErrorKind::InvalidData,
46            "peer frame length outside limits",
47        ));
48    }
49    let mut bytes = Vec::new();
50    bytes.try_reserve_exact(length).map_err(io::Error::other)?;
51    bytes.resize(length, 0);
52    io.read_exact(&mut bytes).await?;
53    let mut trailing = [0];
54    if io.read(&mut trailing).await? != 0 {
55        return Err(io::Error::new(
56            io::ErrorKind::InvalidData,
57            "trailing peer frame bytes",
58        ));
59    }
60    let value = serde_json::from_slice(&bytes)
61        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
62    Ok((value, length))
63}
64
65struct BoundedVec {
66    bytes: Vec<u8>,
67    limit: usize,
68}
69
70impl Write for BoundedVec {
71    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
72        if bytes.len() > self.limit - self.bytes.len() {
73            return Err(io::Error::new(
74                io::ErrorKind::InvalidInput,
75                "encoded peer frame exceeds byte limit",
76            ));
77        }
78        let needed = self.bytes.len() + bytes.len();
79        if needed > self.bytes.capacity() {
80            // Amortize small serializer writes, without reserving past the cap.
81            let capacity = needed
82                .max(self.bytes.capacity().saturating_mul(2))
83                .min(self.limit);
84            self.bytes
85                .try_reserve_exact(capacity - self.bytes.len())
86                .map_err(io::Error::other)?;
87        }
88        self.bytes.extend_from_slice(bytes);
89        Ok(bytes.len())
90    }
91    fn flush(&mut self) -> io::Result<()> {
92        Ok(())
93    }
94}
95
96async fn write<T: AsyncWrite + Unpin + Send, V: Serialize>(
97    io: &mut T,
98    value: &V,
99    limit: usize,
100) -> io::Result<()> {
101    let mut body = BoundedVec {
102        bytes: Vec::new(),
103        limit,
104    };
105    serde_json::to_writer(&mut body, value)
106        .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
107    let length = u32::try_from(body.bytes.len()).map_err(io::Error::other)?;
108    io.write_all(&length.to_be_bytes()).await?;
109    io.write_all(&body.bytes).await?;
110    io.close().await
111}
112
113#[async_trait]
114impl libp2p::request_response::Codec for Codec {
115    type Protocol = StreamProtocol;
116    type Request = Request;
117    type Response = Response;
118
119    async fn read_request<T>(
120        &mut self,
121        negotiated: &StreamProtocol,
122        io: &mut T,
123    ) -> io::Result<Request>
124    where
125        T: AsyncRead + Unpin + Send,
126    {
127        protocol(negotiated)?;
128        Ok(read(io, REQUEST_LIMIT).await?.0)
129    }
130    /// The reader does not know the request kind, so a frame of up to
131    /// RESPONSE_LIMIT is read first; the 256 KiB inventory bound then applies.
132    async fn read_response<T>(
133        &mut self,
134        negotiated: &StreamProtocol,
135        io: &mut T,
136    ) -> io::Result<Response>
137    where
138        T: AsyncRead + Unpin + Send,
139    {
140        protocol(negotiated)?;
141        let (response, length) = read(io, RESPONSE_LIMIT).await?;
142        if length > response_limit(&response) {
143            return Err(io::Error::new(
144                io::ErrorKind::InvalidData,
145                "peer inventory or gossip frame exceeds its limit",
146            ));
147        }
148        Ok(response)
149    }
150    async fn write_request<T>(
151        &mut self,
152        negotiated: &StreamProtocol,
153        io: &mut T,
154        request: Request,
155    ) -> io::Result<()>
156    where
157        T: AsyncWrite + Unpin + Send,
158    {
159        protocol(negotiated)?;
160        write(io, &request, REQUEST_LIMIT).await
161    }
162    async fn write_response<T>(
163        &mut self,
164        negotiated: &StreamProtocol,
165        io: &mut T,
166        response: Response,
167    ) -> io::Result<()>
168    where
169        T: AsyncWrite + Unpin + Send,
170    {
171        protocol(negotiated)?;
172        let limit = response_limit(&response);
173        write(io, &response, limit).await
174    }
175}