use std::borrow::Cow;
use std::collections::HashMap;
use crate::error::{ClientError, MemcacheError};
use super::meta_command::{MetaCommand, MetaOp, MetaResponse, ReturnCode, base64_decode, encode_key};
const OPAQUE_MAX: usize = 32;
fn invalid<T>(message: &'static str) -> Result<T, MemcacheError> {
Err(ClientError::Error(Cow::Borrowed(message)).into())
}
fn opaque_flag(opaque: &[u8]) -> Result<Vec<u8>, MemcacheError> {
if opaque.is_empty() || opaque.len() > OPAQUE_MAX {
return invalid("opaque must be 1-32 bytes");
}
if opaque.iter().any(|&byte| byte <= 0x20 || byte == 0x7f) {
return invalid("opaque must not contain whitespace or control bytes");
}
let mut flag = Vec::with_capacity(opaque.len() + 1);
flag.push(b'O');
flag.extend_from_slice(opaque);
Ok(flag)
}
struct FlagBuilder(Vec<Vec<u8>>);
impl FlagBuilder {
fn new() -> FlagBuilder {
FlagBuilder(Vec::new())
}
fn marker(&mut self, enabled: bool, wire: &'static [u8]) {
if enabled {
self.0.push(wire.to_vec());
}
}
fn token(&mut self, prefix: u8, value: Option<u64>) {
if let Some(value) = value {
let mut flag = vec![prefix];
flag.extend_from_slice(value.to_string().as_bytes());
self.0.push(flag);
}
}
fn opaque(&mut self, opaque: Option<&[u8]>) -> Result<(), MemcacheError> {
if let Some(opaque) = opaque {
self.0.push(opaque_flag(opaque)?);
}
Ok(())
}
fn into_inner(self) -> Vec<Vec<u8>> {
self.0
}
}
#[derive(Debug, Clone)]
pub struct GetOptions {
pub value: bool,
pub return_client_flags: bool,
pub return_cas: bool,
pub return_ttl: bool,
pub return_size: bool,
pub return_last_access: bool,
pub return_hit_before: bool,
pub return_key: bool,
pub no_lru_bump: bool,
pub touch: Option<u32>,
pub vivify_ttl: Option<u32>,
pub recache_ttl: Option<u32>,
pub unless_cas: Option<u64>,
pub new_cas: Option<u64>,
pub opaque: Option<Vec<u8>>,
}
impl Default for GetOptions {
fn default() -> GetOptions {
GetOptions {
value: true,
return_client_flags: false,
return_cas: false,
return_ttl: false,
return_size: false,
return_last_access: false,
return_hit_before: false,
return_key: false,
no_lru_bump: false,
touch: None,
vivify_ttl: None,
recache_ttl: None,
unless_cas: None,
new_cas: None,
opaque: None,
}
}
}
pub fn build_get(key: impl Into<Vec<u8>>, options: &GetOptions) -> Result<MetaCommand, MemcacheError> {
if options.recache_ttl == Some(0) {
return invalid("recache_ttl must be >= 1");
}
let mut flags = FlagBuilder::new();
flags.marker(options.value, b"v");
flags.marker(options.return_client_flags, b"f");
flags.marker(options.return_cas, b"c");
flags.marker(options.return_ttl, b"t");
flags.marker(options.return_size, b"s");
flags.marker(options.return_last_access, b"l");
flags.marker(options.return_hit_before, b"h");
flags.marker(options.return_key, b"k");
flags.marker(options.no_lru_bump, b"u");
flags.token(b'T', options.touch.map(u64::from));
flags.token(b'N', options.vivify_ttl.map(u64::from));
flags.token(b'R', options.recache_ttl.map(u64::from));
flags.token(b'C', options.unless_cas);
flags.token(b'E', options.new_cas);
flags.opaque(options.opaque.as_deref())?;
let mut command = MetaCommand::new(MetaOp::Get, key);
command.flags = flags.into_inner();
Ok(command)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SetMode {
#[default]
Set,
Add,
Replace,
Append,
Prepend,
}
impl SetMode {
fn flag(self) -> Option<&'static [u8]> {
match self {
SetMode::Set => None,
SetMode::Add => Some(b"ME"),
SetMode::Replace => Some(b"MR"),
SetMode::Append => Some(b"MA"),
SetMode::Prepend => Some(b"MP"),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct SetOptions {
pub client_flags: Option<u32>,
pub ttl: Option<u32>,
pub mode: SetMode,
pub compare_cas: Option<u64>,
pub new_cas: Option<u64>,
pub invalidate: bool,
pub vivify_ttl: Option<u32>,
pub return_cas: bool,
pub return_size: bool,
pub return_key: bool,
pub opaque: Option<Vec<u8>>,
}
pub fn build_set(
key: impl Into<Vec<u8>>,
value: impl Into<Vec<u8>>,
options: &SetOptions,
) -> Result<MetaCommand, MemcacheError> {
let concatenation = matches!(options.mode, SetMode::Append | SetMode::Prepend);
if options.vivify_ttl.is_some() && !concatenation {
return invalid("vivify_ttl is only valid for append/prepend");
}
if options.ttl.is_some() && concatenation {
return invalid("ttl is ignored for append/prepend; use vivify_ttl");
}
if options.compare_cas.is_some() && options.mode == SetMode::Add {
return invalid("compare_cas cannot be combined with add mode");
}
let mut flags = FlagBuilder::new();
if let Some(mode_flag) = options.mode.flag() {
flags.marker(true, mode_flag);
}
flags.marker(options.invalidate, b"I");
flags.marker(options.return_cas, b"c");
flags.marker(options.return_size, b"s");
flags.marker(options.return_key, b"k");
flags.token(b'F', options.client_flags.map(u64::from));
flags.token(b'T', options.ttl.map(u64::from));
flags.token(b'C', options.compare_cas);
flags.token(b'E', options.new_cas);
flags.token(b'N', options.vivify_ttl.map(u64::from));
flags.opaque(options.opaque.as_deref())?;
let mut command = MetaCommand::new(MetaOp::Set, key);
command.flags = flags.into_inner();
command.value = Some(value.into());
Ok(command)
}
#[derive(Debug, Clone, Default)]
pub struct DeleteOptions {
pub compare_cas: Option<u64>,
pub new_cas: Option<u64>,
pub invalidate: bool,
pub ttl: Option<u32>,
pub drop_value: bool,
pub return_key: bool,
pub opaque: Option<Vec<u8>>,
}
pub fn build_delete(key: impl Into<Vec<u8>>, options: &DeleteOptions) -> Result<MetaCommand, MemcacheError> {
if options.ttl.is_some() && !options.invalidate {
return invalid("ttl is only applied when invalidate is set");
}
let mut flags = FlagBuilder::new();
flags.marker(options.invalidate, b"I");
flags.marker(options.drop_value, b"x");
flags.marker(options.return_key, b"k");
flags.token(b'C', options.compare_cas);
flags.token(b'E', options.new_cas);
flags.token(b'T', options.ttl.map(u64::from));
flags.opaque(options.opaque.as_deref())?;
let mut command = MetaCommand::new(MetaOp::Delete, key);
command.flags = flags.into_inner();
Ok(command)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ArithmeticMode {
#[default]
Increment,
Decrement,
}
#[derive(Debug, Clone)]
pub struct ArithmeticOptions {
pub delta: Option<u64>,
pub mode: ArithmeticMode,
pub initial: Option<u64>,
pub initial_ttl: Option<u32>,
pub ttl: Option<u32>,
pub compare_cas: Option<u64>,
pub new_cas: Option<u64>,
pub return_value: bool,
pub return_ttl: bool,
pub return_cas: bool,
pub return_key: bool,
pub opaque: Option<Vec<u8>>,
}
impl Default for ArithmeticOptions {
fn default() -> ArithmeticOptions {
ArithmeticOptions {
delta: None,
mode: ArithmeticMode::Increment,
initial: None,
initial_ttl: None,
ttl: None,
compare_cas: None,
new_cas: None,
return_value: true,
return_ttl: false,
return_cas: false,
return_key: false,
opaque: None,
}
}
}
pub fn build_arithmetic(key: impl Into<Vec<u8>>, options: &ArithmeticOptions) -> Result<MetaCommand, MemcacheError> {
if options.initial.is_some() && options.initial_ttl.is_none() {
return invalid("initial requires initial_ttl to vivify on miss");
}
let mut flags = FlagBuilder::new();
flags.marker(options.mode == ArithmeticMode::Decrement, b"MD");
flags.marker(options.return_value, b"v");
flags.marker(options.return_ttl, b"t");
flags.marker(options.return_cas, b"c");
flags.marker(options.return_key, b"k");
flags.token(b'D', options.delta);
flags.token(b'N', options.initial_ttl.map(u64::from));
flags.token(b'J', options.initial);
flags.token(b'T', options.ttl.map(u64::from));
flags.token(b'C', options.compare_cas);
flags.token(b'E', options.new_cas);
flags.opaque(options.opaque.as_deref())?;
let mut command = MetaCommand::new(MetaOp::Arithmetic, key);
command.flags = flags.into_inner();
Ok(command)
}
pub fn build_noop() -> MetaCommand {
MetaCommand::new(MetaOp::Noop, Vec::new())
}
pub fn build_debug(key: impl Into<Vec<u8>>) -> Result<MetaCommand, MemcacheError> {
let key = key.into();
let (_, needs_base64) = encode_key(&key)?;
if needs_base64 {
return invalid("meta debug does not support keys that require base64");
}
Ok(MetaCommand::new(MetaOp::Debug, key))
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MetaCommandResult {
pub rc: ReturnCode,
pub value: Option<Vec<u8>>,
pub cas: Option<u64>,
pub ttl: Option<i64>,
pub client_flags: Option<u32>,
pub size: Option<u64>,
pub last_access: Option<u64>,
pub hit_before: Option<bool>,
pub key: Option<Vec<u8>>,
pub opaque: Option<Vec<u8>>,
pub won: bool,
pub busy: bool,
pub stale: bool,
pub flags: Vec<Vec<u8>>,
}
impl MetaCommandResult {
pub fn ok(&self) -> bool {
matches!(self.rc, ReturnCode::Hd | ReturnCode::Va)
}
}
fn parse_int<T: std::str::FromStr<Err = std::num::ParseIntError>>(token: &[u8]) -> Result<T, MemcacheError> {
Ok(std::str::from_utf8(token)?.parse::<T>()?)
}
pub fn parse_meta_result(response: MetaResponse) -> Result<MetaCommandResult, MemcacheError> {
let mut result = MetaCommandResult {
rc: response.rc,
value: response.value,
cas: None,
ttl: None,
client_flags: None,
size: None,
last_access: None,
hit_before: None,
key: None,
opaque: None,
won: false,
busy: false,
stale: false,
flags: response.flags,
};
let mut key_base64 = false;
for flag in &result.flags {
let Some((&code, token)) = flag.split_first() else {
continue;
};
match code {
b'f' => result.client_flags = Some(parse_int(token)?),
b'c' => result.cas = Some(parse_int(token)?),
b't' => result.ttl = Some(parse_int(token)?),
b'l' => result.last_access = Some(parse_int(token)?),
b's' => result.size = Some(parse_int(token)?),
b'h' => result.hit_before = Some(token != b"0"),
b'k' => result.key = Some(token.to_vec()),
b'O' => result.opaque = Some(token.to_vec()),
b'W' => result.won = true,
b'Z' => result.busy = true,
b'X' => result.stale = true,
b'b' => key_base64 = true,
_ => {}
}
}
if key_base64 && let Some(key) = &result.key {
result.key = Some(base64_decode(key)?);
}
Ok(result)
}
pub fn parse_debug_result(response: &MetaResponse) -> Result<Option<HashMap<String, String>>, MemcacheError> {
if response.rc == ReturnCode::En {
return Ok(None);
}
if response.rc != ReturnCode::Me {
return Err(crate::error::ServerError::BadResponse(Cow::Owned(format!(
"unexpected debug response {:?}",
response.rc
)))
.into());
}
let mut fields = HashMap::new();
for token in response.flags.iter().skip(1) {
let mut split = token.splitn(2, |&byte| byte == b'=');
let name = split.next().unwrap_or_default();
let value = split.next().unwrap_or_default();
fields.insert(String::from_utf8(name.to_vec())?, String::from_utf8(value.to_vec())?);
}
Ok(Some(fields))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_get_default() {
let command = build_get("foo", &GetOptions::default()).unwrap();
assert_eq!(command.encode().unwrap(), b"mg foo v\r\n".to_vec());
}
#[test]
fn build_get_all_flags() {
let options = GetOptions {
value: true,
return_client_flags: true,
return_cas: true,
return_ttl: true,
return_size: true,
return_last_access: true,
return_hit_before: true,
return_key: true,
no_lru_bump: true,
touch: Some(60),
vivify_ttl: Some(30),
recache_ttl: Some(10),
unless_cas: Some(7),
new_cas: Some(8),
opaque: Some(b"tok".to_vec()),
};
let command = build_get("foo", &options).unwrap();
assert_eq!(
command.encode().unwrap(),
b"mg foo v f c t s l h k u T60 N30 R10 C7 E8 Otok\r\n".to_vec()
);
}
#[test]
fn build_get_validation() {
let options = GetOptions {
recache_ttl: Some(0),
..GetOptions::default()
};
assert!(build_get("foo", &options).is_err());
}
#[test]
fn build_set_modes() {
let command = build_set("foo", "bar", &SetOptions::default()).unwrap();
assert_eq!(command.encode().unwrap(), b"ms foo 3\r\nbar\r\n".to_vec());
let options = SetOptions {
mode: SetMode::Add,
ttl: Some(60),
client_flags: Some(1),
return_cas: true,
..SetOptions::default()
};
let command = build_set("foo", "bar", &options).unwrap();
assert_eq!(command.encode().unwrap(), b"ms foo 3 ME c F1 T60\r\nbar\r\n".to_vec());
}
#[test]
fn build_set_validation() {
let append_with_ttl = SetOptions {
mode: SetMode::Append,
ttl: Some(60),
..SetOptions::default()
};
assert!(build_set("foo", "bar", &append_with_ttl).is_err());
let set_with_vivify = SetOptions {
vivify_ttl: Some(60),
..SetOptions::default()
};
assert!(build_set("foo", "bar", &set_with_vivify).is_err());
let add_with_cas = SetOptions {
mode: SetMode::Add,
compare_cas: Some(1),
..SetOptions::default()
};
assert!(build_set("foo", "bar", &add_with_cas).is_err());
}
#[test]
fn build_delete_flags() {
let command = build_delete("foo", &DeleteOptions::default()).unwrap();
assert_eq!(command.encode().unwrap(), b"md foo\r\n".to_vec());
let options = DeleteOptions {
invalidate: true,
ttl: Some(30),
compare_cas: Some(5),
..DeleteOptions::default()
};
let command = build_delete("foo", &options).unwrap();
assert_eq!(command.encode().unwrap(), b"md foo I C5 T30\r\n".to_vec());
let ttl_without_invalidate = DeleteOptions {
ttl: Some(30),
..DeleteOptions::default()
};
assert!(build_delete("foo", &ttl_without_invalidate).is_err());
}
#[test]
fn build_arithmetic_flags() {
let command = build_arithmetic("counter", &ArithmeticOptions::default()).unwrap();
assert_eq!(command.encode().unwrap(), b"ma counter v\r\n".to_vec());
let options = ArithmeticOptions {
mode: ArithmeticMode::Decrement,
delta: Some(2),
initial: Some(0),
initial_ttl: Some(60),
..ArithmeticOptions::default()
};
let command = build_arithmetic("counter", &options).unwrap();
assert_eq!(command.encode().unwrap(), b"ma counter MD v D2 N60 J0\r\n".to_vec());
let initial_without_ttl = ArithmeticOptions {
initial: Some(0),
..ArithmeticOptions::default()
};
assert!(build_arithmetic("counter", &initial_without_ttl).is_err());
}
#[test]
fn build_noop_and_debug() {
assert_eq!(build_noop().encode().unwrap(), b"mn\r\n".to_vec());
assert_eq!(build_debug("foo").unwrap().encode().unwrap(), b"me foo\r\n".to_vec());
assert!(build_debug(b"a key".to_vec()).is_err());
}
#[test]
fn opaque_validation() {
for opaque in [&b""[..], &[b'x'; 33][..], b"a b", b"a\x7f"] {
let options = GetOptions {
opaque: Some(opaque.to_vec()),
..GetOptions::default()
};
assert!(build_get("foo", &options).is_err(), "opaque {:?} should fail", opaque);
}
}
#[test]
fn parse_meta_result_flags() {
let mut response = MetaResponse::parse_header(b"VA 3 f1 c42 t-1 h1 l5 s3 W Otok kZm9v b").unwrap();
response.value = Some(b"bar".to_vec());
let result = parse_meta_result(response).unwrap();
assert!(result.ok());
assert_eq!(result.rc, ReturnCode::Va);
assert_eq!(result.value, Some(b"bar".to_vec()));
assert_eq!(result.client_flags, Some(1));
assert_eq!(result.cas, Some(42));
assert_eq!(result.ttl, Some(-1));
assert_eq!(result.hit_before, Some(true));
assert_eq!(result.last_access, Some(5));
assert_eq!(result.size, Some(3));
assert!(result.won);
assert!(!result.busy);
assert!(!result.stale);
assert_eq!(result.opaque, Some(b"tok".to_vec()));
assert_eq!(result.key, Some(b"foo".to_vec()));
}
#[test]
fn parse_meta_result_miss() {
let response = MetaResponse::parse_header(b"EN").unwrap();
let result = parse_meta_result(response).unwrap();
assert!(!result.ok());
assert_eq!(result.rc, ReturnCode::En);
}
#[test]
fn parse_debug_result_fields() {
let response = MetaResponse::parse_header(b"ME foo exp=-1 la=2 cas=3").unwrap();
let fields = parse_debug_result(&response).unwrap().unwrap();
assert_eq!(fields.get("exp").map(String::as_str), Some("-1"));
assert_eq!(fields.get("la").map(String::as_str), Some("2"));
assert_eq!(fields.get("cas").map(String::as_str), Some("3"));
let miss = MetaResponse::parse_header(b"EN").unwrap();
assert!(parse_debug_result(&miss).unwrap().is_none());
let unexpected = MetaResponse::parse_header(b"HD").unwrap();
assert!(parse_debug_result(&unexpected).is_err());
}
}