use serde_json::Value;
pub const MAGIC: &[u8; 4] = b"ZBXD";
pub const FLAG_PROTOCOL: u8 = 0x01;
pub const FLAG_COMPRESSED: u8 = 0x02;
pub const FLAG_LARGE: u8 = 0x04;
pub const HEADER_LEN: usize = 13;
pub const LARGE_HEADER_LEN: usize = 21;
pub const MAX_DATA_BYTES: usize = 1024 * 1024;
pub const MAX_ITEMS: usize = 1000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HeaderError {
BadMagic,
BadFlags(u8),
Compressed,
TooLarge(u64),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
pub flags: u8,
pub data_len: u64,
pub reserved: u64,
}
pub fn header_len(flags: u8) -> usize {
if flags & FLAG_LARGE != 0 {
LARGE_HEADER_LEN
} else {
HEADER_LEN
}
}
pub fn parse_header(bytes: &[u8]) -> Result<Header, HeaderError> {
if bytes.len() < 5 || &bytes[..4] != MAGIC {
return Err(HeaderError::BadMagic);
}
let flags = bytes[4];
if flags & FLAG_PROTOCOL == 0 || flags & !(FLAG_PROTOCOL | FLAG_COMPRESSED | FLAG_LARGE) != 0 {
return Err(HeaderError::BadFlags(flags));
}
let need = header_len(flags);
if bytes.len() < need {
return Err(HeaderError::BadMagic);
}
let (data_len, reserved) = if flags & FLAG_LARGE != 0 {
(
u64::from_le_bytes(bytes[5..13].try_into().unwrap()),
u64::from_le_bytes(bytes[13..21].try_into().unwrap()),
)
} else {
(
u64::from(u32::from_le_bytes(bytes[5..9].try_into().unwrap())),
u64::from(u32::from_le_bytes(bytes[9..13].try_into().unwrap())),
)
};
if flags & FLAG_COMPRESSED != 0 {
return Err(HeaderError::Compressed);
}
if data_len > MAX_DATA_BYTES as u64 {
return Err(HeaderError::TooLarge(data_len));
}
Ok(Header {
flags,
data_len,
reserved,
})
}
pub fn encode(payload: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(HEADER_LEN + payload.len());
out.extend_from_slice(MAGIC);
out.push(FLAG_PROTOCOL);
out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
out.extend_from_slice(&0u32.to_le_bytes());
out.extend_from_slice(payload);
out
}
pub fn encode_large(payload: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(LARGE_HEADER_LEN + payload.len());
out.extend_from_slice(MAGIC);
out.push(FLAG_PROTOCOL | FLAG_LARGE);
out.extend_from_slice(&(payload.len() as u64).to_le_bytes());
out.extend_from_slice(&0u64.to_le_bytes());
out.extend_from_slice(payload);
out
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SenderItem {
pub host: String,
pub key: String,
pub value: String,
pub clock: Option<i64>,
pub ns: Option<i64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Request {
SenderData {
items: Vec<SenderItem>,
clock: Option<i64>,
},
Other(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestError {
NotJson,
NoRequest,
BadData,
TooManyItems,
}
impl RequestError {
pub fn info(self) -> &'static str {
match self {
RequestError::NotJson => "cannot parse request as a JSON object",
RequestError::NoRequest => "cannot find the \"request\" tag",
RequestError::BadData => "cannot parse the \"data\" array",
RequestError::TooManyItems => "too many values in one request",
}
}
}
fn text(v: Option<&Value>) -> Option<String> {
match v? {
Value::String(s) => Some(s.clone()),
Value::Number(n) => Some(n.to_string()),
Value::Bool(b) => Some(b.to_string()),
_ => None,
}
}
fn int(v: Option<&Value>) -> Option<i64> {
match v? {
Value::Number(n) => n.as_i64(),
Value::String(s) => s.trim().parse().ok(),
_ => None,
}
}
pub fn parse_request(data: &[u8]) -> Result<Request, RequestError> {
let json: Value = serde_json::from_slice(data).map_err(|_| RequestError::NotJson)?;
let obj = json.as_object().ok_or(RequestError::NotJson)?;
let request = obj
.get("request")
.and_then(Value::as_str)
.ok_or(RequestError::NoRequest)?;
if request != "sender data" {
return Ok(Request::Other(request.to_string()));
}
let data = obj
.get("data")
.and_then(Value::as_array)
.ok_or(RequestError::BadData)?;
if data.len() > MAX_ITEMS {
return Err(RequestError::TooManyItems);
}
let mut items = Vec::with_capacity(data.len());
for entry in data {
let entry = entry.as_object().ok_or(RequestError::BadData)?;
items.push(SenderItem {
host: text(entry.get("host")).ok_or(RequestError::BadData)?,
key: text(entry.get("key")).ok_or(RequestError::BadData)?,
value: text(entry.get("value")).unwrap_or_default(),
clock: int(entry.get("clock")),
ns: int(entry.get("ns")),
});
}
Ok(Request::SenderData {
items,
clock: int(obj.get("clock")),
})
}
pub fn info_string(processed: u64, failed: u64, total: u64, seconds: f64) -> String {
format!("processed: {processed}; failed: {failed}; total: {total}; seconds spent: {seconds:.6}")
}
pub fn render_result(processed: u64, failed: u64, total: u64, seconds: f64) -> Vec<u8> {
let body = serde_json::json!({
"response": "success",
"info": info_string(processed, failed, total, seconds),
});
encode(body.to_string().as_bytes())
}
pub fn render_failed(info: &'static str) -> Vec<u8> {
let body = serde_json::json!({"response": "failed", "info": info});
encode(body.to_string().as_bytes())
}
pub fn read_result(packet: &[u8]) -> Option<(u64, u64, u64)> {
let header = parse_header(packet).ok()?;
let body = packet.get(header_len(header.flags)..)?;
let json: Value = serde_json::from_slice(body).ok()?;
if json.get("response")?.as_str()? != "success" {
return None;
}
let info = json.get("info")?.as_str()?;
let mut fields = info.split("; ");
let mut next = |name: &str| -> Option<u64> {
fields
.next()?
.strip_prefix(name)?
.strip_prefix(": ")?
.parse()
.ok()
};
Some((next("processed")?, next("failed")?, next("total")?))
}