netget/server/zeromq/
wire.rs

1//! ZMTP 3.1 (ZeroMQ's wire protocol, RFC 37) with the NULL mechanism: the 64-byte greeting,
2//! the READY handshake with Socket-Type and Identity metadata, multipart frames, and the
3//! PING/PONG, SUBSCRIBE/CANCEL and ERROR commands. Shared by the server and the client.
4use anyhow::{bail, ensure, Context, Result};
5use std::time::Duration;
6use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
7
8pub const GREETING_LEN: usize = 64;
9/// Bytes in one frame, and in one whole multipart message.
10pub const MAX_MESSAGE_BYTES: usize = 1024 * 1024;
11/// Frames in one multipart message.
12pub const MAX_FRAMES: usize = 64;
13/// Bytes in one command (READY metadata, SUBSCRIBE topics, PING contexts).
14pub const MAX_COMMAND_BYTES: usize = 64 * 1024;
15/// Deadline for the greeting and READY exchange, and for each write.
16pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
17
18const MORE: u8 = 0x01;
19const LONG: u8 = 0x02;
20const COMMAND: u8 = 0x04;
21
22/// The greeting NetGet sends: ZMTP 3.1, NULL mechanism, as-server false (NULL has no roles).
23pub fn greeting() -> [u8; GREETING_LEN] {
24    let mut g = [0u8; GREETING_LEN];
25    g[0] = 0xFF;
26    g[9] = 0x7F;
27    g[10] = 3;
28    g[11] = 1;
29    g[12..16].copy_from_slice(b"NULL");
30    g
31}
32
33/// Check a peer's greeting; returns its (major, minor) version.
34pub fn check_greeting(g: &[u8; GREETING_LEN]) -> Result<(u8, u8)> {
35    ensure!(g[0] == 0xFF && g[9] & 0x01 == 0x01, "not a ZMTP 3 greeting");
36    ensure!(
37        g[10] >= 3,
38        "ZMTP {}.{} is not supported (3.0 or later)",
39        g[10],
40        g[11]
41    );
42    let mechanism: Vec<u8> = g[12..32].iter().copied().take_while(|b| *b != 0).collect();
43    ensure!(
44        mechanism == b"NULL",
45        "security mechanism {:?} is not supported (NULL only)",
46        String::from_utf8_lossy(&mechanism)
47    );
48    Ok((g[10], g[11]))
49}
50
51/// One frame on the wire.
52pub fn encode_frame(body: &[u8], more: bool, command: bool) -> Vec<u8> {
53    let mut flags = 0u8;
54    if more {
55        flags |= MORE;
56    }
57    if command {
58        flags |= COMMAND;
59    }
60    let mut out = Vec::with_capacity(body.len() + 9);
61    if body.len() > 255 {
62        out.push(flags | LONG);
63        out.extend_from_slice(&(body.len() as u64).to_be_bytes());
64    } else {
65        out.push(flags);
66        out.push(body.len() as u8);
67    }
68    out.extend_from_slice(body);
69    out
70}
71
72/// A multipart message: every frame but the last carries MORE.
73pub fn encode_message(frames: &[Vec<u8>]) -> Vec<u8> {
74    let mut out = Vec::new();
75    for (i, frame) in frames.iter().enumerate() {
76        out.extend(encode_frame(frame, i + 1 < frames.len(), false));
77    }
78    if frames.is_empty() {
79        out.extend(encode_frame(&[], false, false));
80    }
81    out
82}
83
84/// A command frame: name, then its body.
85pub fn encode_command(name: &str, body: &[u8]) -> Vec<u8> {
86    let mut b = Vec::with_capacity(1 + name.len() + body.len());
87    b.push(name.len() as u8);
88    b.extend_from_slice(name.as_bytes());
89    b.extend_from_slice(body);
90    encode_frame(&b, false, true)
91}
92
93/// READY with Socket-Type and, when given, Identity.
94pub fn ready(socket_type: &str, identity: Option<&[u8]>) -> Vec<u8> {
95    let mut props = Vec::new();
96    let mut put = |name: &str, value: &[u8]| {
97        props.push(name.len() as u8);
98        props.extend_from_slice(name.as_bytes());
99        props.extend_from_slice(&(value.len() as u32).to_be_bytes());
100        props.extend_from_slice(value);
101    };
102    put("Socket-Type", socket_type.as_bytes());
103    if let Some(id) = identity {
104        put("Identity", id);
105    }
106    encode_command("READY", &props)
107}
108
109/// ERROR with a reason (at most 255 bytes).
110pub fn error(reason: &str) -> Vec<u8> {
111    let reason = &reason.as_bytes()[..reason.len().min(255)];
112    let mut body = vec![reason.len() as u8];
113    body.extend_from_slice(reason);
114    encode_command("ERROR", &body)
115}
116
117/// Parse READY metadata into (name, value) pairs.
118pub fn parse_metadata(mut body: &[u8]) -> Result<Vec<(String, Vec<u8>)>> {
119    let mut out = Vec::new();
120    while !body.is_empty() {
121        let n = body[0] as usize;
122        ensure!(body.len() >= 1 + n + 4, "truncated READY property");
123        let name = String::from_utf8_lossy(&body[1..1 + n]).to_string();
124        let len = u32::from_be_bytes(body[1 + n..5 + n].try_into()?) as usize;
125        ensure!(body.len() >= 5 + n + len, "truncated READY property value");
126        out.push((name, body[5 + n..5 + n + len].to_vec()));
127        body = &body[5 + n + len..];
128    }
129    Ok(out)
130}
131
132/// Whether two socket types may talk (ZMTP 3.1's compatibility table).
133pub fn compatible(ours: &str, theirs: &str) -> bool {
134    matches!(
135        (ours, theirs),
136        ("REQ", "REP" | "ROUTER")
137            | ("REP", "REQ" | "DEALER")
138            | ("DEALER", "REP" | "DEALER" | "ROUTER")
139            | ("ROUTER", "REQ" | "DEALER" | "ROUTER")
140            | ("PUB", "SUB" | "XSUB")
141            | ("XPUB", "SUB" | "XSUB")
142            | ("SUB", "PUB" | "XPUB")
143            | ("XSUB", "PUB" | "XPUB")
144            | ("PUSH", "PULL")
145            | ("PULL", "PUSH")
146            | ("PAIR", "PAIR")
147    )
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
151pub enum Incoming {
152    Message(Vec<Vec<u8>>),
153    /// A command other than READY: its name and body.
154    Command(String, Vec<u8>),
155}
156
157async fn read_frame<R: AsyncRead + Unpin>(r: &mut R, max: usize) -> Result<(u8, Vec<u8>)> {
158    let flags = r.read_u8().await?;
159    ensure!(
160        flags & !(MORE | LONG | COMMAND) == 0,
161        "reserved ZMTP frame flags 0x{flags:02x}"
162    );
163    let size = if flags & LONG != 0 {
164        r.read_u64().await?
165    } else {
166        u64::from(r.read_u8().await?)
167    };
168    ensure!(
169        size <= max as u64,
170        "ZMTP frame of {size} bytes exceeds {max}"
171    );
172    let mut body = vec![0u8; size as usize];
173    r.read_exact(&mut body).await?;
174    Ok((flags, body))
175}
176
177/// Read one whole message or command. `None` on a clean EOF between them; the first byte may
178/// wait `idle`, the rest of the message must follow within `HANDSHAKE_TIMEOUT` per frame.
179pub async fn read_incoming<R: AsyncRead + Unpin>(
180    r: &mut R,
181    idle: Duration,
182) -> Result<Option<Incoming>> {
183    let mut first = [0u8; 1];
184    match tokio::time::timeout(idle, r.read(&mut first)).await {
185        Ok(Ok(0)) => return Ok(None),
186        Ok(Ok(_)) => {}
187        Ok(Err(e)) => return Err(e.into()),
188        Err(_) => bail!("ZMTP peer idle for {}s", idle.as_secs()),
189    }
190    let mut chained = (&first[..]).chain(&mut *r);
191    let mut frames = Vec::new();
192    let mut total = 0usize;
193    loop {
194        let (flags, body) = tokio::time::timeout(
195            HANDSHAKE_TIMEOUT,
196            read_frame(&mut chained, MAX_MESSAGE_BYTES - total),
197        )
198        .await
199        .context("ZMTP frame deadline")??;
200        if flags & COMMAND != 0 {
201            ensure!(frames.is_empty(), "ZMTP command inside a multipart message");
202            ensure!(body.len() <= MAX_COMMAND_BYTES, "ZMTP command too large");
203            ensure!(
204                !body.is_empty() && body.len() > body[0] as usize,
205                "truncated ZMTP command"
206            );
207            let n = body[0] as usize;
208            let name = String::from_utf8_lossy(&body[1..1 + n]).to_string();
209            return Ok(Some(Incoming::Command(name, body[1 + n..].to_vec())));
210        }
211        total += body.len();
212        frames.push(body);
213        ensure!(
214            frames.len() <= MAX_FRAMES,
215            "ZMTP message of more than {MAX_FRAMES} frames"
216        );
217        if flags & MORE == 0 {
218            return Ok(Some(Incoming::Message(frames)));
219        }
220    }
221}
222
223/// The greeting and READY exchange. Returns the peer's (socket type, identity).
224pub async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
225    stream: &mut S,
226    socket_type: &str,
227    identity: Option<&[u8]>,
228) -> Result<(String, Vec<u8>)> {
229    tokio::time::timeout(HANDSHAKE_TIMEOUT, async {
230        stream.write_all(&greeting()).await?;
231        let mut theirs = [0u8; GREETING_LEN];
232        stream.read_exact(&mut theirs).await?;
233        if let Err(e) = check_greeting(&theirs) {
234            let _ = stream.write_all(&error(&e.to_string())).await;
235            return Err(e);
236        }
237        stream.write_all(&ready(socket_type, identity)).await?;
238        let (flags, body) = read_frame(stream, MAX_COMMAND_BYTES).await?;
239        ensure!(
240            flags & COMMAND != 0,
241            "ZMTP peer sent a message before READY"
242        );
243        ensure!(body.len() > 5 && body[0] == 5, "malformed ZMTP command");
244        match &body[1..6] {
245            b"READY" => {}
246            b"ERROR" => bail!(
247                "ZMTP peer refused: {}",
248                String::from_utf8_lossy(&body[7.min(body.len())..])
249            ),
250            other => bail!(
251                "ZMTP peer sent {} before READY",
252                String::from_utf8_lossy(other)
253            ),
254        }
255        let props = parse_metadata(&body[6..])?;
256        let mut peer_type = String::new();
257        let mut peer_identity = Vec::new();
258        for (name, value) in props {
259            match name.to_ascii_lowercase().as_str() {
260                "socket-type" => peer_type = String::from_utf8_lossy(&value).to_string(),
261                "identity" => peer_identity = value,
262                _ => {}
263            }
264        }
265        if !compatible(socket_type, &peer_type) {
266            let reason = format!("{socket_type} cannot talk to {peer_type}");
267            let _ = stream.write_all(&error(&reason)).await;
268            bail!(reason);
269        }
270        Ok((peer_type, peer_identity))
271    })
272    .await
273    .context("ZMTP handshake deadline")?
274}
275
276/// Frames as event text: each UTF-8 frame as itself, and hex for the whole message when any
277/// frame is not UTF-8.
278pub fn frames_to_text(frames: &[Vec<u8>]) -> (Vec<String>, &'static str) {
279    if frames.iter().all(|f| std::str::from_utf8(f).is_ok()) {
280        (
281            frames
282                .iter()
283                .map(|f| String::from_utf8_lossy(f).to_string())
284                .collect(),
285            "utf8",
286        )
287    } else {
288        (
289            frames
290                .iter()
291                .map(|f| f.iter().map(|b| format!("{b:02x}")).collect())
292                .collect(),
293            "hex",
294        )
295    }
296}
297
298/// Frames from an action, by its declared encoding.
299pub fn frames_from_text(
300    frames: &[serde_json::Value],
301    encoding: Option<&str>,
302) -> Result<Vec<Vec<u8>>> {
303    ensure!(frames.len() <= MAX_FRAMES, "at most {MAX_FRAMES} frames");
304    let mut out = Vec::with_capacity(frames.len());
305    let mut total = 0;
306    for f in frames {
307        let text = f.as_str().context("each frame must be a string")?;
308        let bytes = match encoding.unwrap_or("utf8") {
309            "utf8" => text.as_bytes().to_vec(),
310            "hex" => {
311                ensure!(
312                    text.is_ascii() && text.len().is_multiple_of(2),
313                    "hex frames must be an even number of hex digits"
314                );
315                (0..text.len())
316                    .step_by(2)
317                    .map(|i| u8::from_str_radix(&text[i..i + 2], 16).context("invalid hex frame"))
318                    .collect::<Result<Vec<u8>>>()?
319            }
320            other => bail!("encoding must be utf8 or hex, not {other}"),
321        };
322        total += bytes.len();
323        out.push(bytes);
324    }
325    ensure!(
326        total <= MAX_MESSAGE_BYTES,
327        "message exceeds {MAX_MESSAGE_BYTES} bytes"
328    );
329    Ok(out)
330}
netget/server/zeromq/wire.rs

netget/server/zeromq/
wire.rs
