1use 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}
220pub 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}