1use 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
27fn 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 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 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}