Skip to main content

dig_pex/
wire.rs

1//! The PEX wire — the four `type`-tagged JSON messages (SPEC §4) and their framing.
2//!
3//! ## The four messages
4//!
5//! | `type` | Purpose |
6//! |---|---|
7//! | `pex_handshake` | first message each direction: version + network + declared interval + own flags (§4.2) |
8//! | `pex_snapshot` | the first data message: a capped picture of the sender's first-hand set (§4.3) |
9//! | `pex_delta` | the periodic message: `added` / `dropped` relative to what this link was told (§4.4) |
10//! | `pex_error` | the advisory error envelope (§4.5) |
11//!
12//! ## Framing (SPEC §4.1)
13//!
14//! - **Byte stream** (node↔node, §10.1): [`PexMessage::encode`] / [`PexMessage::decode`] frame each
15//!   message as a **`u32` big-endian length prefix + JSON body** — byte-identical to the dig-nat /
16//!   dig-dht wires. A length prefix over [`crate::caps::PEX_MAX_FRAME`] is rejected
17//!   without allocating the body.
18//! - **Relay WebSocket** (relay→node, §10.2): [`PexMessage::to_json`] / [`PexMessage::from_json`]
19//!   carry the bare JSON object in one WebSocket text frame (the socket already delimits messages).
20//!
21//! JSON shapes are **frozen**: the `type` tags and field names are the wire contract, and unknown
22//! fields are ignored on receive (additive evolution).
23
24use 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/// A PEX protocol message — one of the four `type`-tagged shapes (SPEC §4). The `type` tag and field
33/// names are the frozen wire contract.
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(tag = "type", rename_all = "snake_case")]
36pub enum PexMessage {
37    /// The first PEX message a participant sends on a link, in each direction, before anything else
38    /// (SPEC §4.2). Declares the wire `version`, the sender's `network_id`, its minimum send
39    /// `interval` (seconds), and its own capability `flags`.
40    PexHandshake {
41        /// The PEX wire version the sender speaks (this crate: `1`).
42        version: u32,
43        /// The sender's network. MUST match the receiver's.
44        network_id: String,
45        /// Seconds — the sender's declared minimum spacing between its own data messages. MUST be
46        /// within `[30, 3600]`; a receiver clamps out-of-range values for enforcement.
47        interval: u32,
48        /// The sender's own capability flags (optional; defaults to empty).
49        #[serde(default)]
50        flags: Vec<String>,
51    },
52    /// The first **data message** in a direction — a fuller, capped picture of the sender's
53    /// first-hand known-peer set, so a fresh link warms up in one message (SPEC §4.3). Exactly one
54    /// per direction per link.
55    PexSnapshot {
56        /// The sender's first-hand peers (MAY be empty; at most `PEX_MAX_SNAPSHOT`, freshest-first).
57        peers: Vec<PeerEntry>,
58    },
59    /// The periodic message: what changed in the sender's first-hand set **relative to what this link
60    /// has already been told** (SPEC §4.4). A delta with both arrays empty MUST NOT be sent.
61    PexDelta {
62        /// Newly-known entries, or already-told entries whose advertised content changed (an update
63        /// replaces the previous entry). At most `PEX_MAX_ADDED`.
64        added: Vec<PeerEntry>,
65        /// `<64hex>` peer ids the sender no longer considers good (SPEC §8.3). At most
66        /// `PEX_MAX_DROPPED`. Advisory — a receiver never deletes a first-hand-verified peer on it.
67        dropped: Vec<String>,
68    },
69    /// The advisory error envelope, either direction (SPEC §4.5). Best-effort; never requires a reply.
70    PexError {
71        /// A stable error code (SPEC §4.5 table).
72        code: u16,
73        /// A human-readable message.
74        message: String,
75    },
76}
77
78impl PexMessage {
79    /// Serialize as a `u32` big-endian length prefix + JSON body — the byte-stream framing (§4.1).
80    #[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    /// Read + decode one framed message from `r`, bounded by [`PEX_MAX_FRAME`]. A length prefix over
90    /// the bound errors **before** allocating or reading the body (SPEC §4.1, §7.2).
91    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    /// The bare JSON object as a string — the relay WebSocket text-frame form (§4.1, §10.2).
107    #[must_use]
108    pub fn to_json(&self) -> String {
109        serde_json::to_string(self).expect("pex message serializes")
110    }
111
112    /// The bare JSON object as bytes.
113    #[must_use]
114    pub fn to_json_bytes(&self) -> Vec<u8> {
115        serde_json::to_vec(self).expect("pex message serializes")
116    }
117
118    /// Parse a bare JSON object (the relay WebSocket text-frame form). An `Err` here is a
119    /// `PEX_BAD_MESSAGE` (SPEC §7.3) the caller strikes.
120    pub fn from_json(s: &str) -> Result<Self, serde_json::Error> {
121        serde_json::from_str(s)
122    }
123
124    /// The message's `type` tag string (`"pex_handshake"` | `"pex_snapshot"` | `"pex_delta"` |
125    /// `"pex_error"`).
126    #[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        // u32-BE length prefix.
238        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        // Additive evolution: a future field must not break a v1 decoder (SPEC §2, §4.1).
265        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        // A snapshot missing `peers` is a malformed message (SPEC §7.3), not an empty snapshot.
278        assert!(PexMessage::from_json(r#"{"type":"pex_snapshot"}"#).is_err());
279    }
280}