netget/client/beanstalkd/
wire.rs

1//! Structured Beanstalkd requests and bounded, correlated responses.
2use crate::server::beanstalkd::wire::{self, Command};
3use anyhow::{bail, ensure, Context, Result};
4use serde_json::{json, Value};
5use std::time::Duration;
6use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt};
7pub const DEADLINE: Duration = Duration::from_secs(30);
8pub const MAX_REPLY_BYTES: usize = 1024 * 1024;
9#[derive(Debug)]
10pub struct Request {
11    pub command: Command,
12    pub bytes: Vec<u8>,
13}
14impl Request {
15    pub fn from_action(v: &Value) -> Result<Self> {
16        let op = v["operation"].as_str().context("Missing operation")?;
17        let number = |key: &str, default: Option<u64>| -> Result<u64> {
18            v.get(key)
19                .map(|v| {
20                    v.as_u64()
21                        .with_context(|| format!("{key} must be an unsigned integer"))
22                })
23                .unwrap_or_else(|| default.with_context(|| format!("Missing {key}")))
24        };
25        let tube = || -> Result<&str> {
26            let tube = v["tube"].as_str().context("Missing tube")?;
27            ensure!(wire::valid_tube_name(tube), "Invalid tube name");
28            Ok(tube)
29        };
30        let id = || -> Result<u64> {
31            let id = number("id", None)?;
32            ensure!(id > 0, "Job id must be positive");
33            Ok(id)
34        };
35        let priority = || number("priority", Some(1024));
36        let delay = || number("delay", Some(0));
37        let mut body = None;
38        let line = match op {
39            "put" => {
40                let text = v["body"].as_str().context("Missing text body")?;
41                ensure!(
42                    text.len() <= wire::MAX_JOB_BYTES,
43                    "Job body exceeds 65535 bytes"
44                );
45                body = Some(text);
46                format!(
47                    "put {} {} {} {}",
48                    priority()?,
49                    delay()?,
50                    number("ttr", Some(60))?,
51                    text.len()
52                )
53            }
54            "use" | "watch" | "ignore" => format!("{op} {}", tube()?),
55            "reserve" => {
56                let seconds = number("timeout_secs", Some(0))?;
57                ensure!(seconds <= 25, "Reservation timeout must be 0..=25 seconds");
58                format!("reserve-with-timeout {seconds}")
59            }
60            "reserve_job" | "delete" | "touch" | "peek" | "kick_job" | "stats_job" => {
61                format!("{} {}", op.replace('_', "-"), id()?)
62            }
63            "release" => format!("release {} {} {}", id()?, priority()?, delay()?),
64            "bury" => format!("bury {} {}", id()?, priority()?),
65            "kick" => format!("kick {}", number("count", Some(1))?),
66            "stats_tube" => format!("stats-tube {}", tube()?),
67            "pause_tube" => format!("pause-tube {} {}", tube()?, delay()?),
68            "stats" | "list_tubes" | "list_tube_used" | "list_tubes_watched" | "peek_ready"
69            | "peek_delayed" | "peek_buried" | "quit" => op.replace('_', "-"),
70            _ => bail!("Unsupported Beanstalkd operation"),
71        };
72        ensure!(
73            line.len() + 2 <= wire::MAX_LINE_BYTES,
74            "Command header too long"
75        );
76        let command = wire::parse_command(&line);
77        ensure!(
78            !matches!(command, Command::BadFormat | Command::Unknown),
79            "Invalid command arguments or numeric range"
80        );
81        let mut bytes = format!("{line}\r\n").into_bytes();
82        if let Some(body) = body {
83            bytes.extend_from_slice(body.as_bytes());
84            bytes.extend_from_slice(b"\r\n");
85        }
86        Ok(Self { command, bytes })
87    }
88}
89fn unsigned(text: &str) -> Result<u64> {
90    ensure!(
91        !text.is_empty() && text.bytes().all(|c| c.is_ascii_digit()),
92        "Invalid unsigned response number"
93    );
94    Ok(text.parse()?)
95}
96fn job_id(text: &str) -> Result<u64> {
97    let id = unsigned(text)?;
98    ensure!(id > 0, "Job ID must be positive");
99    Ok(id)
100}
101async fn line<R: AsyncBufRead + Unpin>(r: &mut R) -> Result<String> {
102    let mut bytes = Vec::new();
103    loop {
104        let chunk = r.fill_buf().await?;
105        ensure!(!chunk.is_empty(), "Truncated Beanstalkd response");
106        let n = chunk
107            .iter()
108            .position(|b| *b == b'\n')
109            .map(|n| n + 1)
110            .unwrap_or(chunk.len());
111        ensure!(
112            bytes.len() + n <= wire::MAX_LINE_BYTES,
113            "Beanstalkd response header too long"
114        );
115        bytes.extend_from_slice(&chunk[..n]);
116        r.consume(n);
117        if bytes.ends_with(b"\n") {
118            ensure!(bytes.ends_with(b"\r\n"), "Beanstalkd requires CRLF");
119            bytes.truncate(bytes.len() - 2);
120            return Ok(String::from_utf8(bytes)?);
121        }
122    }
123}
124pub async fn response<R: AsyncBufRead + Unpin>(reader: &mut R, request: &Request) -> Result<Value> {
125    tokio::time::timeout(DEADLINE, async {
126        let header = line(reader).await?;
127        let args: Vec<&str> = header.split(' ').collect();
128        let status = args.first().copied().context("Missing status")?;
129        let mut wire = format!("{header}\r\n").into_bytes();
130        let mut payload = None;
131        if matches!(status, "RESERVED" | "FOUND" | "OK") {
132            ensure!(
133                args.len() == if status == "OK" { 2 } else { 3 },
134                "Malformed byte-counted response header"
135            );
136            let length = usize::try_from(unsigned(args.last().context("Missing byte count")?)?)?;
137            ensure!(
138                length <= MAX_REPLY_BYTES,
139                "Beanstalkd reply body exceeds limit"
140            );
141            if status != "OK" {
142                ensure!(
143                    length <= wire::MAX_JOB_BYTES,
144                    "Beanstalkd job exceeds limit"
145                );
146            }
147            let mut body = vec![0; length];
148            reader.read_exact(&mut body).await?;
149            let mut ending = [0; 2];
150            reader.read_exact(&mut ending).await?;
151            ensure!(ending == *b"\r\n", "Body missing CRLF");
152            wire.extend_from_slice(&body);
153            wire.extend_from_slice(&ending);
154            payload = Some(
155                String::from_utf8(body)
156                    .context("Non-UTF8 jobs are outside this client's text payload scope")?,
157            );
158        }
159        let global_error = matches!(
160            status,
161            "OUT_OF_MEMORY"
162                | "INTERNAL_ERROR"
163                | "BAD_FORMAT"
164                | "UNKNOWN_COMMAND"
165                | "JOB_TOO_BIG"
166                | "EXPECTED_CRLF"
167        ) && args.len() == 1;
168        let fits = match &request.command {
169            Command::Use(tube) => args == vec!["USING", tube.as_str()],
170            Command::Watch(_) => status == "WATCHING" && args.len() == 2,
171            Command::Ignore(_) => {
172                status == "NOT_IGNORED" && args.len() == 1
173                    || status == "WATCHING" && args.len() == 2
174            }
175            Command::ListTubeUsed => {
176                status == "USING" && args.len() == 2 && wire::valid_tube_name(args[1])
177            }
178            Command::ListTubesWatched => status == "OK" && args.len() == 2,
179            _ => wire::reply_fits(&request.command, &wire),
180        };
181        ensure!(
182            global_error || fits,
183            "Unexpected Beanstalkd response for request"
184        );
185        let mut result = json!({"status":status});
186        if status == "OK" {
187            let value = parse_yaml(
188                payload.as_deref().context("Missing YAML payload")?,
189                matches!(
190                    request.command,
191                    Command::ListTubes | Command::ListTubesWatched
192                ),
193            )?;
194            match request.command {
195                Command::ListTubes | Command::ListTubesWatched => ensure!(
196                    value
197                        .as_array()
198                        .is_some_and(|v| v.iter().all(Value::is_string)),
199                    "Expected tube-name list"
200                ),
201                _ => ensure!(value.is_object(), "Expected stats mapping"),
202            }
203            result["data"] = value;
204        } else if matches!(status, "RESERVED" | "FOUND") {
205            ensure!(args.len() == 3, "Malformed job response");
206            let id = job_id(args[1])?;
207            if let Command::Peek(expected) | Command::ReserveJob(expected) = request.command {
208                ensure!(id == expected, "Reply job ID does not match request");
209            }
210            result["id"] = json!(id);
211            result["body"] = json!(payload);
212        } else if args.len() == 2 {
213            match status {
214                "USING" => result["tube"] = json!(args[1]),
215                "WATCHING" | "KICKED" => {
216                    let count = unsigned(args[1])?;
217                    ensure!(
218                        status != "WATCHING" || count > 0,
219                        "Empty watch set is invalid"
220                    );
221                    result["count"] = json!(count);
222                }
223                "INSERTED" | "BURIED" => result["id"] = json!(job_id(args[1])?),
224                _ => bail!("Unknown status arguments"),
225            }
226        }
227        Ok(result)
228    })
229    .await
230    .context("Beanstalkd response deadline exceeded")?
231}
232
233/// Deserialize scalar values without ever visiting a sequence/map's children.
234/// Deserializing into Value first would allocate nested graphs before rejecting them.
235struct Scalar(Value);
236impl<'de> serde::Deserialize<'de> for Scalar {
237    fn deserialize<D: serde::Deserializer<'de>>(
238        deserializer: D,
239    ) -> std::result::Result<Self, D::Error> {
240        struct ScalarVisitor;
241        impl<'de> serde::de::Visitor<'de> for ScalarVisitor {
242            type Value = Scalar;
243            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244                f.write_str("a flat YAML scalar")
245            }
246            fn visit_bool<E: serde::de::Error>(self, v: bool) -> std::result::Result<Scalar, E> {
247                Ok(Scalar(json!(v)))
248            }
249            fn visit_i64<E: serde::de::Error>(self, v: i64) -> std::result::Result<Scalar, E> {
250                Ok(Scalar(json!(v)))
251            }
252            fn visit_u64<E: serde::de::Error>(self, v: u64) -> std::result::Result<Scalar, E> {
253                Ok(Scalar(json!(v)))
254            }
255            fn visit_f64<E: serde::de::Error>(self, v: f64) -> std::result::Result<Scalar, E> {
256                serde_json::Number::from_f64(v)
257                    .map(|n| Scalar(Value::Number(n)))
258                    .ok_or_else(|| E::custom("Nonfinite YAML number"))
259            }
260            fn visit_str<E: serde::de::Error>(self, v: &str) -> std::result::Result<Scalar, E> {
261                Ok(Scalar(json!(v)))
262            }
263            fn visit_string<E: serde::de::Error>(
264                self,
265                v: String,
266            ) -> std::result::Result<Scalar, E> {
267                Ok(Scalar(Value::String(v)))
268            }
269            fn visit_unit<E: serde::de::Error>(self) -> std::result::Result<Scalar, E> {
270                Ok(Scalar(Value::Null))
271            }
272            // The default visit_seq, visit_map and visit_enum reject immediately.
273        }
274        deserializer.deserialize_any(ScalarVisitor)
275    }
276}
277
278/// Beanstalkd's YAML payload is a flat list or scalar mapping. Parse each scalar
279/// independently so aliases/anchors cannot expand across records.
280pub fn parse_yaml(text: &str, list: bool) -> Result<Value> {
281    let text = text
282        .strip_prefix("---\n")
283        .context("Missing YAML document marker")?;
284    let rows: Vec<&str> = text.lines().collect();
285    ensure!(
286        rows.len() <= wire::MAX_YAML_ENTRIES,
287        "Too many YAML entries"
288    );
289    if list {
290        let mut out = Vec::new();
291        for row in rows {
292            if row.trim().is_empty() {
293                continue;
294            }
295            let text = row.strip_prefix("- ").context("Expected flat tube list")?;
296            let value = if wire::valid_tube_name(text) {
297                json!(text)
298            } else {
299                let Scalar(value) = serde_yaml::from_str(text)?;
300                value
301            };
302            let name = value.as_str().context("Tube name must be text")?;
303            ensure!(wire::valid_tube_name(name), "Invalid tube name in list");
304            out.push(value);
305        }
306        Ok(Value::Array(out))
307    } else {
308        let mut out = serde_json::Map::new();
309        for row in rows {
310            if row.trim().is_empty() {
311                continue;
312            }
313            let (key, text) = row
314                .split_once(": ")
315                .context("Expected flat stats mapping")?;
316            ensure!(
317                !key.is_empty()
318                    && key
319                        .bytes()
320                        .all(|c| c.is_ascii_alphanumeric() || c == b'-' || c == b'_'),
321                "Invalid stats key"
322            );
323            // v1.12 emits uname fields without YAML quoting: os starts with '#'
324            // on Linux and contains ': ' on Darwin. Treat only these protocol text
325            // fields as opaque strings, never deserialize arbitrary nested values.
326            let value = if matches!(key, "hostname" | "os" | "platform")
327                && !text.starts_with(['\"', '\''])
328            {
329                json!(text)
330            } else if matches!(key, "tube" | "name") && wire::valid_tube_name(text) {
331                json!(text)
332            } else {
333                let Scalar(value) = serde_yaml::from_str(text)?;
334                value
335            };
336            ensure!(
337                out.insert(key.into(), value).is_none(),
338                "Duplicate stats key"
339            );
340        }
341        Ok(Value::Object(out))
342    }
343}
netget/client/beanstalkd/wire.rs

netget/client/beanstalkd/
wire.rs
