Skip to main content

netget/client/gearman/
wire.rs

1//! Typed Gearman binary requests and a continuous bounded response reader.
2use crate::server::gearman::wire as gear;
3use anyhow::{bail, Context, Result};
4use serde_json::Value;
5use std::time::Duration;
6use tokio::io::{AsyncRead, AsyncReadExt};
7pub const DEADLINE: Duration = Duration::from_secs(15);
8pub const MAX_FOLLOWUPS: u8 = 4;
9pub const MAX_JOBS: usize = 64;
10pub const MAX_ABILITIES: usize = 64;
11pub const NOOP: u32 = 6;
12pub const NO_JOB: u32 = 10;
13pub const JOB_ASSIGN: u32 = 11;
14pub const JOB_ASSIGN_UNIQ: u32 = 31;
15pub const DEFAULT_ROLE: &str = "submitter";
16
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum Role {
19    Submitter,
20    Worker,
21}
22impl Role {
23    pub fn parse(role: &str) -> Result<Self> {
24        match role {
25            "submitter" => Ok(Self::Submitter),
26            "worker" => Ok(Self::Worker),
27            _ => bail!("Gearman role must be submitter or worker"),
28        }
29    }
30    pub fn name(self) -> &'static str {
31        match self {
32            Self::Submitter => "submitter",
33            Self::Worker => "worker",
34        }
35    }
36}
37#[derive(Debug)]
38pub struct Request {
39    pub packet_type: u32,
40    pub args: Vec<Vec<u8>>,
41}
42impl Request {
43    pub fn bytes(&self) -> Vec<u8> {
44        gear::encode(
45            true,
46            self.packet_type,
47            &self.args.iter().map(Vec::as_slice).collect::<Vec<_>>(),
48        )
49    }
50    pub fn expects_reply(&self) -> bool {
51        matches!(
52            self.packet_type,
53            gear::SUBMIT_JOB
54                | gear::SUBMIT_JOB_HIGH
55                | gear::SUBMIT_JOB_LOW
56                | gear::SUBMIT_JOB_BG
57                | gear::SUBMIT_JOB_HIGH_BG
58                | gear::SUBMIT_JOB_LOW_BG
59                | gear::GET_STATUS
60                | gear::ECHO_REQ
61                | gear::OPTION_REQ
62                | gear::GRAB_JOB
63                | gear::GRAB_JOB_UNIQ
64        )
65    }
66    pub fn worker_only(&self) -> bool {
67        gear::is_worker_packet(self.packet_type) || self.packet_type == gear::SET_CLIENT_ID
68    }
69    pub fn submit(&self) -> bool {
70        matches!(
71            self.packet_type,
72            gear::SUBMIT_JOB
73                | gear::SUBMIT_JOB_HIGH
74                | gear::SUBMIT_JOB_LOW
75                | gear::SUBMIT_JOB_BG
76                | gear::SUBMIT_JOB_HIGH_BG
77                | gear::SUBMIT_JOB_LOW_BG
78        )
79    }
80    pub fn background(&self) -> bool {
81        matches!(
82            self.packet_type,
83            gear::SUBMIT_JOB_BG | gear::SUBMIT_JOB_HIGH_BG | gear::SUBMIT_JOB_LOW_BG
84        )
85    }
86}
87fn text<'a>(v: &'a Value, key: &str) -> Result<&'a str> {
88    v[key]
89        .as_str()
90        .with_context(|| format!("Gearman {key} must be a string"))
91}
92fn field(s: &str, key: &str, min: usize, max: usize) -> Result<Vec<u8>> {
93    anyhow::ensure!(
94        (min..=max).contains(&s.len()) && !s.contains('\0'),
95        "Gearman {key} must have {min}..={max} UTF-8 bytes without NUL"
96    );
97    Ok(s.as_bytes().to_vec())
98}
99pub fn handle(bytes: &[u8]) -> Result<String> {
100    anyhow::ensure!(
101        !bytes.is_empty() && bytes.len() < 64 && bytes.iter().all(|b| b.is_ascii_graphic()),
102        "Gearman job handle must have 1..63 printable ASCII bytes"
103    );
104    Ok(std::str::from_utf8(bytes)?.into())
105}
106fn job(v: &Value) -> Result<Vec<u8>> {
107    let h = text(v, "job_handle")?;
108    handle(h.as_bytes())?;
109    Ok(h.as_bytes().to_vec())
110}
111fn function(v: &Value) -> Result<Vec<u8>> {
112    field(
113        text(v, "function_name")?,
114        "function",
115        1,
116        gear::MAX_FUNCTION_NAME,
117    )
118}
119pub fn request(v: &Value) -> Result<Request> {
120    let operation = text(v, "operation")?;
121    let (packet_type, args) = match v["type"].as_str() {
122        Some("gearman_request") => match operation {
123            "submit" => {
124                let priority = match v.get("priority") {
125                    None => "normal",
126                    Some(p) => p.as_str().context("Gearman priority must be a string")?,
127                };
128                let background = match v.get("background") {
129                    None => false,
130                    Some(b) => b.as_bool().context("Gearman background must be boolean")?,
131                };
132                let t = match (priority, background) {
133                    ("normal", false) => gear::SUBMIT_JOB,
134                    ("normal", true) => gear::SUBMIT_JOB_BG,
135                    ("high", false) => gear::SUBMIT_JOB_HIGH,
136                    ("high", true) => gear::SUBMIT_JOB_HIGH_BG,
137                    ("low", false) => gear::SUBMIT_JOB_LOW,
138                    ("low", true) => gear::SUBMIT_JOB_LOW_BG,
139                    _ => bail!("Gearman priority must be normal, high or low"),
140                };
141                let unique = match v.get("unique_id") {
142                    None => "",
143                    Some(s) => s.as_str().context("Gearman unique_id must be a string")?,
144                };
145                (
146                    t,
147                    vec![
148                        function(v)?,
149                        field(unique, "unique_id", 0, gear::MAX_UNIQUE)?,
150                        text(v, "workload")?.as_bytes().to_vec(),
151                    ],
152                )
153            }
154            "status" => (gear::GET_STATUS, vec![job(v)?]),
155            "echo" => (gear::ECHO_REQ, vec![text(v, "data")?.as_bytes().to_vec()]),
156            "enable_exceptions" => (gear::OPTION_REQ, vec![b"exceptions".to_vec()]),
157            _ => bail!("Unknown Gearman submitter operation"),
158        },
159        Some("gearman_worker") => match operation {
160            "register" => (gear::CAN_DO, vec![function(v)?]),
161            "unregister" => (gear::CANT_DO, vec![function(v)?]),
162            "reset" => (gear::RESET_ABILITIES, vec![]),
163            "grab" => (gear::GRAB_JOB, vec![]),
164            "grab_unique" => (gear::GRAB_JOB_UNIQ, vec![]),
165            "sleep" => (gear::PRE_SLEEP, vec![]),
166            "set_id" => (
167                gear::SET_CLIENT_ID,
168                vec![field(text(v, "client_id")?, "client_id", 1, 64)?],
169            ),
170            "progress" => {
171                let n = v["numerator"]
172                    .as_u64()
173                    .context("Gearman numerator must be a nonnegative integer")?;
174                let d = v["denominator"]
175                    .as_u64()
176                    .context("Gearman denominator must be a positive integer")?;
177                anyhow::ensure!(d>0&&n<=d&&d<=u32::MAX as u64,"Gearman progress must satisfy 0 <= numerator <= denominator <= u32::MAX and denominator > 0");
178                (
179                    gear::WORK_STATUS,
180                    vec![
181                        job(v)?,
182                        n.to_string().into_bytes(),
183                        d.to_string().into_bytes(),
184                    ],
185                )
186            }
187            "data" => (
188                gear::WORK_DATA,
189                vec![job(v)?, text(v, "data")?.as_bytes().to_vec()],
190            ),
191            "warning" => (
192                gear::WORK_WARNING,
193                vec![job(v)?, text(v, "data")?.as_bytes().to_vec()],
194            ),
195            "complete" => (
196                gear::WORK_COMPLETE,
197                vec![job(v)?, text(v, "result")?.as_bytes().to_vec()],
198            ),
199            "fail" => (gear::WORK_FAIL, vec![job(v)?]),
200            "exception" => (
201                gear::WORK_EXCEPTION,
202                vec![job(v)?, text(v, "text")?.as_bytes().to_vec()],
203            ),
204            _ => bail!("Unknown Gearman worker operation"),
205        },
206        _ => bail!("Unknown Gearman client action"),
207    };
208    let size = args.iter().map(Vec::len).sum::<usize>() + args.len().saturating_sub(1);
209    anyhow::ensure!(
210        size <= gear::MAX_PACKET_BYTES,
211        "Gearman packet body exceeds 1 MiB"
212    );
213    Ok(Request { packet_type, args })
214}
215#[derive(Debug)]
216pub struct Frame {
217    pub packet_type: u32,
218    pub data: Vec<u8>,
219}
220/// A dedicated reader owns each complete frame. Idle workers may wait indefinitely;
221/// partial headers/bodies get a deadline and the size is checked before allocation.
222pub async fn frame<R: AsyncRead + Unpin>(r: &mut R) -> Result<Option<Frame>> {
223    let mut header = [0; gear::HEADER_LEN];
224    if r.read(&mut header[..1]).await? == 0 {
225        return Ok(None);
226    }
227    tokio::time::timeout(DEADLINE, async {
228        r.read_exact(&mut header[1..]).await?;
229        let h = gear::parse_header(&header)
230            .map_err(|e| anyhow::anyhow!("Invalid Gearman response header: {e:?}"))?;
231        anyhow::ensure!(!h.request, "Gearman peer sent request magic");
232        let mut data = vec![0; h.size as usize];
233        r.read_exact(&mut data).await?;
234        Ok(Some(Frame {
235            packet_type: h.packet_type,
236            data,
237        }))
238    })
239    .await
240    .context("Gearman partial frame deadline")?
241}
242pub fn args(data: &[u8], n: usize) -> Result<Vec<&[u8]>> {
243    gear::split_args(data, n).context("Malformed Gearman response arguments")
244}
245pub fn number(data: &[u8]) -> Result<u64> {
246    anyhow::ensure!(
247        !data.is_empty() && data.len() <= 10 && data.iter().all(u8::is_ascii_digit),
248        "Invalid Gearman response integer"
249    );
250    let n = std::str::from_utf8(data)?.parse::<u64>()?;
251    anyhow::ensure!(n <= u32::MAX as u64, "Gearman response integer exceeds u32");
252    Ok(n)
253}
254pub fn body(data: &[u8]) -> Value {
255    serde_json::json!({"text":std::str::from_utf8(data).ok(),"bytes":data.len(),"utf8":std::str::from_utf8(data).is_ok()})
256}