1use std::io;
25
26use serde::{Deserialize, Serialize};
27use tokio::io::{AsyncRead, AsyncReadExt};
28
29use crate::caps::PEX_MAX_FRAME;
30use crate::entry::PeerEntry;
31
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(tag = "type", rename_all = "snake_case")]
36pub enum PexMessage {
37 PexHandshake {
41 version: u32,
43 network_id: String,
45 interval: u32,
48 #[serde(default)]
50 flags: Vec<String>,
51 },
52 PexSnapshot {
56 peers: Vec<PeerEntry>,
58 },
59 PexDelta {
62 added: Vec<PeerEntry>,
65 dropped: Vec<String>,
68 },
69 PexError {
71 code: u16,
73 message: String,
75 },
76}
77
78impl PexMessage {
79 #[must_use]
81 pub fn encode(&self) -> Vec<u8> {
82 let body = self.to_json_bytes();
83 let mut out = Vec::with_capacity(4 + body.len());
84 out.extend_from_slice(&(body.len() as u32).to_be_bytes());
85 out.extend_from_slice(&body);
86 out
87 }
88
89 pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
92 let mut len_buf = [0u8; 4];
93 r.read_exact(&mut len_buf).await?;
94 let len = u32::from_be_bytes(len_buf) as usize;
95 if len > PEX_MAX_FRAME {
96 return Err(io::Error::new(
97 io::ErrorKind::InvalidData,
98 "pex frame too large",
99 ));
100 }
101 let mut body = vec![0u8; len];
102 r.read_exact(&mut body).await?;
103 serde_json::from_slice(&body).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
104 }
105
106 #[must_use]
108 pub fn to_json(&self) -> String {
109 serde_json::to_string(self).expect("pex message serializes")
110 }
111
112 #[must_use]
114 pub fn to_json_bytes(&self) -> Vec<u8> {
115 serde_json::to_vec(self).expect("pex message serializes")
116 }
117
118 pub fn from_json(s: &str) -> Result<Self, serde_json::Error> {
121 serde_json::from_str(s)
122 }
123
124 #[must_use]
127 pub fn type_tag(&self) -> &'static str {
128 match self {
129 PexMessage::PexHandshake { .. } => "pex_handshake",
130 PexMessage::PexSnapshot { .. } => "pex_snapshot",
131 PexMessage::PexDelta { .. } => "pex_delta",
132 PexMessage::PexError { .. } => "pex_error",
133 }
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use super::*;
140 use crate::caps::PEX_VERSION;
141 use crate::entry::{Address, PeerEntry, Provenance};
142 use std::io::Cursor;
143
144 fn hex(b: u8) -> String {
145 format!("{b:02x}").repeat(32)
146 }
147
148 fn sample_entry() -> PeerEntry {
149 PeerEntry::new(hex(0x07), "mainnet", 1_719_763_200, Provenance::Direct)
150 .with_address(Address::direct("203.0.113.7", 9444))
151 .with_flag("storage")
152 }
153
154 #[test]
155 fn handshake_frozen_shape() {
156 let m = PexMessage::PexHandshake {
157 version: PEX_VERSION,
158 network_id: "mainnet".into(),
159 interval: 60,
160 flags: vec!["storage".into(), "holepunch".into()],
161 };
162 let s = m.to_json();
163 assert!(s.contains("\"type\":\"pex_handshake\""));
164 assert!(s.contains("\"version\":1"));
165 assert!(s.contains("\"network_id\":\"mainnet\""));
166 assert!(s.contains("\"interval\":60"));
167 assert!(s.contains("\"flags\":[\"storage\",\"holepunch\"]"));
168 }
169
170 #[test]
171 fn snapshot_frozen_shape() {
172 let m = PexMessage::PexSnapshot {
173 peers: vec![sample_entry()],
174 };
175 let s = m.to_json();
176 assert!(s.contains("\"type\":\"pex_snapshot\""));
177 assert!(s.contains("\"peers\":["));
178 }
179
180 #[test]
181 fn delta_frozen_shape() {
182 let m = PexMessage::PexDelta {
183 added: vec![sample_entry()],
184 dropped: vec![hex(0x09)],
185 };
186 let s = m.to_json();
187 assert!(s.contains("\"type\":\"pex_delta\""));
188 assert!(s.contains("\"added\":["));
189 assert!(s.contains("\"dropped\":["));
190 }
191
192 #[test]
193 fn error_frozen_shape() {
194 let m = PexMessage::PexError {
195 code: 3,
196 message: "rate violation".into(),
197 };
198 let s = m.to_json();
199 assert_eq!(
200 s,
201 r#"{"type":"pex_error","code":3,"message":"rate violation"}"#
202 );
203 }
204
205 #[test]
206 fn all_messages_round_trip_through_json() {
207 for m in [
208 PexMessage::PexHandshake {
209 version: 1,
210 network_id: "mainnet".into(),
211 interval: 60,
212 flags: vec!["storage".into()],
213 },
214 PexMessage::PexSnapshot {
215 peers: vec![sample_entry()],
216 },
217 PexMessage::PexDelta {
218 added: vec![sample_entry()],
219 dropped: vec![hex(0x09)],
220 },
221 PexMessage::PexError {
222 code: 6,
223 message: "protocol violation".into(),
224 },
225 ] {
226 let back = PexMessage::from_json(&m.to_json()).unwrap();
227 assert_eq!(m, back);
228 }
229 }
230
231 #[tokio::test]
232 async fn framed_round_trip() {
233 let m = PexMessage::PexSnapshot {
234 peers: vec![sample_entry()],
235 };
236 let bytes = m.encode();
237 let declared = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
239 assert_eq!(declared, bytes.len() - 4);
240 let mut cur = Cursor::new(bytes);
241 let back = PexMessage::decode(&mut cur).await.unwrap();
242 assert_eq!(m, back);
243 }
244
245 #[tokio::test]
246 async fn oversize_length_prefix_rejected_before_body() {
247 let mut buf = ((PEX_MAX_FRAME + 1) as u32).to_be_bytes().to_vec();
248 buf.extend_from_slice(b"{}");
249 let mut cur = Cursor::new(buf);
250 let err = PexMessage::decode(&mut cur).await.unwrap_err();
251 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
252 }
253
254 #[tokio::test]
255 async fn truncated_frame_errors() {
256 let mut buf = 100u32.to_be_bytes().to_vec();
257 buf.extend_from_slice(b"{}");
258 let mut cur = Cursor::new(buf);
259 assert!(PexMessage::decode(&mut cur).await.is_err());
260 }
261
262 #[test]
263 fn unknown_fields_ignored_on_receive() {
264 let s = r#"{"type":"pex_handshake","version":1,"network_id":"mainnet","interval":60,"future":42}"#;
266 let m = PexMessage::from_json(s).unwrap();
267 assert_eq!(m.type_tag(), "pex_handshake");
268 }
269
270 #[test]
271 fn unknown_type_tag_is_error() {
272 assert!(PexMessage::from_json(r#"{"type":"pex_bogus"}"#).is_err());
273 }
274
275 #[test]
276 fn missing_required_field_is_error() {
277 assert!(PexMessage::from_json(r#"{"type":"pex_snapshot"}"#).is_err());
279 }
280}