From 9ffa2528676f42c6344f438615a3cefdb9fd5d18 Mon Sep 17 00:00:00 2001 From: ing-eoking Date: Sat, 24 Oct 2026 22:58:27 +0900 Subject: [PATCH] INTERNAL: Parse vop command lines --- src/ascii.rs | 54 +++-- src/lib.rs | 1 + src/protocol.rs | 526 +++++++++++++++++++++++++++++++++++++++++ src/protocol/tokens.rs | 68 ++++++ 4 files changed, 632 insertions(+), 17 deletions(-) create mode 100644 src/protocol.rs create mode 100644 src/protocol/tokens.rs diff --git a/src/ascii.rs b/src/ascii.rs index 097cd93..5b03e56 100644 --- a/src/ascii.rs +++ b/src/ascii.rs @@ -4,6 +4,7 @@ use std::ptr; use crate::guard; use crate::memcached::ffi::{AUTHZ_NONE, EXTENSION_ASCII_PROTOCOL_DESCRIPTOR, token_t}; +use crate::protocol::{self, Tokens}; const COMMAND: &[u8] = b"vop"; const NOT_SUPPORTED: &[u8] = b"NOT_SUPPORTED\r\n"; @@ -73,12 +74,12 @@ unsafe extern "C" fn accept( /// # Safety /// -/// `response_handler`, if present, must accept `cookie` and bytes that live for the call. +/// `argv`/`argc` must meet [`Tokens::new`]'s contract; `response_handler` must accept `cookie`. unsafe extern "C" fn execute( _cmd_cookie: *const c_void, cookie: *const c_void, - _argc: c_int, - _argv: *mut token_t, + argc: c_int, + argv: *mut token_t, response_handler: ResponseHandler, ) -> bool { guard( @@ -90,14 +91,16 @@ unsafe extern "C" fn execute( let Some(respond) = response_handler else { return false; }; - // SAFETY: `respond` accepts `cookie`, and `NOT_SUPPORTED` is `'static`. - unsafe { - respond( - cookie, - NOT_SUPPORTED.len() as c_int, - NOT_SUPPORTED.as_ptr().cast(), - ) - } + // SAFETY: the caller guarantees `Tokens::new`'s contract. + let tokens = unsafe { Tokens::new(argv, argc) }; + let reply: [&[u8]; 3] = match protocol::parse(&tokens) { + Ok(_) => [NOT_SUPPORTED, b"", b""], + Err(protocol::Error(reason)) => [b"CLIENT_ERROR ", reason.as_bytes(), b"\r\n"], + }; + reply.iter().filter(|part| !part.is_empty()).all(|part| { + // SAFETY: `respond` accepts `cookie`, and every part is `'static`. + unsafe { respond(cookie, part.len() as c_int, part.as_ptr().cast()) } + }) }, ) } @@ -180,21 +183,38 @@ mod tests { true } - #[test] - fn every_vop_line_is_not_supported() { + fn reply_to(line: &str) -> String { let mut reply: Vec = Vec::new(); - let mut tokens = vec![token("vop"), token("create")]; - // SAFETY: `record` gets the unshared `reply` back as its cookie. + let (_buf, mut tokens) = protocol::tokenize(line); + // SAFETY: `tokens` point into `_buf`, and `record` gets the unshared `reply` back. let ok = unsafe { execute( ptr::null(), (&raw mut reply).cast(), - 2, + tokens.len() as c_int, tokens.as_mut_ptr(), Some(record), ) }; assert!(ok); - assert_eq!(reply, b"NOT_SUPPORTED\r\n"); + String::from_utf8(reply).unwrap() + } + + #[test] + fn well_formed_lines_are_not_supported_yet() { + assert_eq!(reply_to("vop create idx 2 0"), "NOT_SUPPORTED\r\n"); + assert_eq!(reply_to("vop drop idx"), "NOT_SUPPORTED\r\n"); + } + + #[test] + fn malformed_lines_get_the_reason() { + assert_eq!( + reply_to("vop create idx"), + "CLIENT_ERROR bad command line format\r\n" + ); + assert_eq!( + reply_to("vop list"), + "CLIENT_ERROR bad command line format\r\n" + ); } } diff --git a/src/lib.rs b/src/lib.rs index 73d66a4..f7bd454 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ mod ascii; mod memcached; +mod protocol; use std::os::raw::c_char; use std::panic::{AssertUnwindSafe, catch_unwind}; diff --git a/src/protocol.rs b/src/protocol.rs new file mode 100644 index 0000000..1094257 --- /dev/null +++ b/src/protocol.rs @@ -0,0 +1,526 @@ +#![cfg_attr( + not(test), + expect( + dead_code, + reason = "the parsed fields are read by the commands built on this parser" + ) +)] + +mod tokens; + +use std::ops::RangeInclusive; +use std::str::FromStr; + +pub use tokens::Tokens; +#[cfg(test)] +pub use tokens::tokenize; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Error(pub &'static str); + +impl Error { + pub const MALFORMED: Self = Self("bad command line format"); + pub const TOO_LARGE: Self = Self("too large value"); +} + +pub const DEFAULT_M: usize = 16; +pub const M_RANGE: RangeInclusive = 2..=1024; +pub const DEFAULT_EFC: usize = 128; +pub const DEFAULT_EFS: usize = 64; +pub const EF_RANGE: RangeInclusive = 1..=65536; +pub const MAX_RESULTS: usize = i32::MAX as usize; +pub const MAX_FILTER_BYTES: usize = 4096; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Metric { + Cos, + L2, + Ip, + Hamming, + Tanimoto, +} + +impl Metric { + fn parse(s: &str) -> Option { + lookup( + s, + &[ + ("cos", Self::Cos), + ("cosine", Self::Cos), + ("l2", Self::L2), + ("l2sq", Self::L2), + ("euclidean", Self::L2), + ("ip", Self::Ip), + ("dot", Self::Ip), + ("hamming", Self::Hamming), + ("tanimoto", Self::Tanimoto), + ("jaccard", Self::Tanimoto), + ], + ) + } + + const fn is_bitwise(self) -> bool { + matches!(self, Self::Hamming | Self::Tanimoto) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Quant { + F32, + F16, + I8, + B1, +} + +impl Quant { + fn parse(s: &str) -> Option { + lookup( + s, + &[ + ("f32", Self::F32), + ("f16", Self::F16), + ("i8", Self::I8), + ("b1", Self::B1), + ], + ) + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum Format { + #[default] + Bin, + Str, +} + +impl Format { + fn parse(s: &str) -> Option { + lookup(s, &[("bin", Self::Bin), ("str", Self::Str)]) + } +} + +/// `vop create [] [metric m] [quant q] [M n] [efc c] [efs s]` +#[derive(Debug)] +pub struct Create<'a> { + pub index: &'a str, + pub dim: u16, + pub exptime: u32, + pub maxcount: Option, + pub metric: Metric, + pub quant: Quant, + pub m: usize, + pub efc: usize, + pub efs: usize, +} + +/// `vop upsert [meta ] [bin|str]` +#[derive(Debug)] +pub struct Upsert<'a> { + pub index: &'a str, + pub vkey: &'a str, + pub vector_bytes: usize, + pub meta_bytes: Option, + pub format: Format, +} + +/// `vop setmeta ` +#[derive(Debug)] +pub struct SetMeta<'a> { + pub index: &'a str, + pub vkey: &'a str, + pub meta_bytes: usize, +} + +#[derive(Debug)] +pub enum Source<'a> { + Vkey(&'a str), + Vector { bytes: usize, format: Format }, +} + +/// `vop search vkey [filter ] [withmeta]`, or +/// `vop search vector [filter ] [withmeta] [bin|str]` +#[derive(Debug)] +pub struct Search<'a> { + pub index: &'a str, + pub source: Source<'a>, + pub k: usize, + pub filter_bytes: Option, + pub with_meta: bool, +} + +#[derive(Debug)] +pub enum Request<'a> { + Create(Create<'a>), + Upsert(Upsert<'a>), + SetMeta(SetMeta<'a>), + Search(Search<'a>), + /// `vop getmeta ` + GetMeta { + index: &'a str, + vkey: &'a str, + }, + /// `vop delete ` + Delete { + index: &'a str, + vkey: &'a str, + }, + /// `vop drop ` + Drop { + index: &'a str, + }, +} + +fn lookup(word: &str, table: &[(&str, T)]) -> Option { + table + .iter() + .find(|(name, _)| word.eq_ignore_ascii_case(name)) + .map(|&(_, value)| value) +} + +fn is(word: &str, keyword: &str) -> bool { + word.eq_ignore_ascii_case(keyword) +} + +/// Sets an option that may be given only once. +fn once(slot: &mut Option, value: T) -> Result<(), Error> { + match slot.replace(value) { + Some(_) => Err(Error::MALFORMED), + None => Ok(()), + } +} + +fn in_range(n: T, range: &RangeInclusive) -> Result { + if range.contains(&n) { + Ok(n) + } else { + Err(Error::MALFORMED) + } +} + +/// Reads the tokens of a line in order. +struct Args<'t, 'a> { + tokens: &'t Tokens<'a>, + at: usize, +} + +impl<'a> Args<'_, 'a> { + fn is_empty(&self) -> bool { + self.at >= self.tokens.len() + } + + fn next(&mut self) -> Result<&'a str, Error> { + let text = self.tokens.text(self.at)?; + self.at += 1; + Ok(text) + } + + fn number(&mut self) -> Result { + self.next()?.parse().map_err(|_| Error::MALFORMED) + } + + /// Reads the next token only if it is a number. + fn optional_number(&mut self) -> Option { + let n = self.tokens.text(self.at).ok()?.parse().ok()?; + self.at += 1; + Some(n) + } +} + +pub fn parse<'a>(tokens: &Tokens<'a>) -> Result, Error> { + let mut args = Args { tokens, at: 1 }; + let sub = args.next()?; + let request = if is(sub, "create") { + Request::Create(create(&mut args)?) + } else if is(sub, "upsert") { + Request::Upsert(upsert(&mut args)?) + } else if is(sub, "setmeta") { + Request::SetMeta(SetMeta { + index: args.next()?, + vkey: args.next()?, + meta_bytes: args.number()?, + }) + } else if is(sub, "search") { + Request::Search(search(&mut args)?) + } else if is(sub, "getmeta") { + Request::GetMeta { + index: args.next()?, + vkey: args.next()?, + } + } else if is(sub, "delete") { + Request::Delete { + index: args.next()?, + vkey: args.next()?, + } + } else if is(sub, "drop") { + Request::Drop { + index: args.next()?, + } + } else { + return Err(Error::MALFORMED); + }; + if !args.is_empty() { + return Err(Error::MALFORMED); + } + Ok(request) +} + +fn create<'a>(args: &mut Args<'_, 'a>) -> Result, Error> { + let index = args.next()?; + let dim = in_range(args.number()?, &(1..=u16::MAX))?; + let exptime = args.number()?; + let maxcount = args + .optional_number() + .map(|n| in_range(n, &(1..=u32::MAX))) + .transpose()?; + + let (mut metric, mut quant, mut m, mut efc, mut efs) = (None, None, None, None, None); + while !args.is_empty() { + let key = args.next()?; + let value = args.next()?; + if is(key, "metric") { + once(&mut metric, Metric::parse(value).ok_or(Error::MALFORMED)?)?; + } else if is(key, "quant") { + once(&mut quant, Quant::parse(value).ok_or(Error::MALFORMED)?)?; + } else if is(key, "M") { + once(&mut m, value.parse().map_err(|_| Error::MALFORMED)?)?; + } else if is(key, "efc") { + once(&mut efc, value.parse().map_err(|_| Error::MALFORMED)?)?; + } else if is(key, "efs") { + once(&mut efs, value.parse().map_err(|_| Error::MALFORMED)?)?; + } else { + return Err(Error::MALFORMED); + } + } + + let metric = metric.unwrap_or(Metric::Cos); + let quant = quant.unwrap_or(Quant::F32); + if metric.is_bitwise() != (quant == Quant::B1) { + return Err(Error::MALFORMED); + } + Ok(Create { + index, + dim, + exptime, + maxcount, + metric, + quant, + m: in_range(m.unwrap_or(DEFAULT_M), &M_RANGE)?, + efc: in_range(efc.unwrap_or(DEFAULT_EFC), &EF_RANGE)?, + efs: in_range(efs.unwrap_or(DEFAULT_EFS), &EF_RANGE)?, + }) +} + +fn upsert<'a>(args: &mut Args<'_, 'a>) -> Result, Error> { + let index = args.next()?; + let vkey = args.next()?; + let vector_bytes = in_range(args.number()?, &(1..=usize::MAX))?; + + let (mut meta_bytes, mut format) = (None, None); + while !args.is_empty() { + let word = args.next()?; + if is(word, "meta") { + once(&mut meta_bytes, args.number()?)?; + } else { + once(&mut format, Format::parse(word).ok_or(Error::MALFORMED)?)?; + } + } + Ok(Upsert { + index, + vkey, + vector_bytes, + meta_bytes, + format: format.unwrap_or_default(), + }) +} + +fn search<'a>(args: &mut Args<'_, 'a>) -> Result, Error> { + let index = args.next()?; + let from = args.next()?; + let (mut source, k) = if is(from, "vkey") { + let vkey = args.next()?; + (Source::Vkey(vkey), args.number()?) + } else if is(from, "vector") { + let k = args.number()?; + let bytes = in_range(args.number()?, &(1..=usize::MAX))?; + let format = Format::default(); + (Source::Vector { bytes, format }, k) + } else { + return Err(Error::MALFORMED); + }; + + let (mut filter_bytes, mut with_meta, mut format) = (None, None, None); + while !args.is_empty() { + let word = args.next()?; + if is(word, "filter") { + let n = args.number()?; + if n > MAX_FILTER_BYTES { + return Err(Error::TOO_LARGE); + } + once(&mut filter_bytes, n)?; + } else if is(word, "withmeta") { + once(&mut with_meta, ())?; + } else if matches!(source, Source::Vector { .. }) { + once(&mut format, Format::parse(word).ok_or(Error::MALFORMED)?)?; + } else { + return Err(Error::MALFORMED); + } + } + if let (Source::Vector { format: slot, .. }, Some(format)) = (&mut source, format) { + *slot = format; + } + + Ok(Search { + index, + source, + k: in_range(k, &(1..=MAX_RESULTS))?, + filter_bytes, + with_meta: with_meta.is_some(), + }) +} + +#[cfg(test)] +mod tests { + use std::os::raw::c_int; + + use super::*; + + fn parse_line(line: &str) -> Result { + let (_buf, raw) = tokenize(line); + // SAFETY: `raw` holds tokens into `_buf`, and both outlive the view. + let tokens = unsafe { Tokens::new(raw.as_ptr(), raw.len() as c_int) }; + parse(&tokens).map(|request| describe(&request)) + } + + fn describe(request: &Request<'_>) -> String { + match request { + Request::Create(c) => format!( + "create {} {} {} {:?} {:?} {:?} {} {} {}", + c.index, c.dim, c.exptime, c.maxcount, c.metric, c.quant, c.m, c.efc, c.efs + ), + Request::Upsert(u) => format!( + "upsert {} {} {} {:?} {:?}", + u.index, u.vkey, u.vector_bytes, u.meta_bytes, u.format + ), + Request::SetMeta(s) => format!("setmeta {} {} {}", s.index, s.vkey, s.meta_bytes), + Request::Search(s) => { + let source = match &s.source { + Source::Vkey(vkey) => format!("vkey {vkey}"), + Source::Vector { bytes, format } => format!("vector {bytes} {format:?}"), + }; + format!( + "search {} {} {} {:?} {}", + s.index, source, s.k, s.filter_bytes, s.with_meta + ) + } + Request::GetMeta { index, vkey } => format!("getmeta {index} {vkey}"), + Request::Delete { index, vkey } => format!("delete {index} {vkey}"), + Request::Drop { index } => format!("drop {index}"), + } + } + + #[test] + fn create_takes_defaults_and_options() { + assert_eq!( + parse_line("vop create idx 128 0").unwrap(), + "create idx 128 0 None Cos F32 16 128 64" + ); + assert_eq!( + parse_line("vop create idx 2 60 100 metric l2 quant f16 M 32 efc 200 efs 10").unwrap(), + "create idx 2 60 Some(100) L2 F16 32 200 10" + ); + assert_eq!( + parse_line("vop CREATE idx 8 0 metric hamming quant b1").unwrap(), + "create idx 8 0 None Hamming B1 16 128 64" + ); + } + + #[test] + fn create_rejects_bad_values() { + for (line, error) in [ + ("vop create idx", Error::MALFORMED), + ("vop create idx 0 0", Error::MALFORMED), + ("vop create idx 65536 0", Error::MALFORMED), + ("vop create idx 2 -1", Error::MALFORMED), + ("vop create idx 2 0 0", Error::MALFORMED), + ("vop create idx 2 0 metric", Error::MALFORMED), + ("vop create idx 2 0 metric foo", Error::MALFORMED), + ("vop create idx 2 0 size 3", Error::MALFORMED), + ("vop create idx 2 0 metric l2 METRIC cos", Error::MALFORMED), + ("vop create idx 2 0 M 1", Error::MALFORMED), + ("vop create idx 2 0 quant b1", Error::MALFORMED), + ] { + assert_eq!(parse_line(line), Err(error), "{line}"); + } + } + + #[test] + fn upsert_and_setmeta_read_their_lengths() { + assert_eq!( + parse_line("vop upsert idx k1 8 meta 5 str").unwrap(), + "upsert idx k1 8 Some(5) Str" + ); + assert_eq!( + parse_line("vop upsert idx k1 8").unwrap(), + "upsert idx k1 8 None Bin" + ); + assert_eq!( + parse_line("vop setmeta idx k1 5").unwrap(), + "setmeta idx k1 5" + ); + assert_eq!( + parse_line("vop upsert idx k1 8 meta 1 meta 2"), + Err(Error::MALFORMED) + ); + assert_eq!( + parse_line("vop upsert idx k1 8 withmeta"), + Err(Error::MALFORMED) + ); + assert_eq!(parse_line("vop setmeta idx k1"), Err(Error::MALFORMED)); + for line in [ + "vop upsert idx k1 0", + "vop upsert idx k1 8 meta", + "vop upsert idx k1 8 bin str", + "vop upsert idx k1 8 str str", + ] { + assert_eq!(parse_line(line), Err(Error::MALFORMED), "{line}"); + } + } + + #[test] + fn search_reads_either_source() { + assert_eq!( + parse_line("vop search idx vkey k1 10 filter 12 withmeta").unwrap(), + "search idx vkey k1 10 Some(12) true" + ); + assert_eq!( + parse_line("vop search idx vector 5 8 str").unwrap(), + "search idx vector 8 Str 5 None false" + ); + assert_eq!( + parse_line("vop search idx vkey k1 0"), + Err(Error::MALFORMED) + ); + assert_eq!( + parse_line("vop search idx vector 5 8 filter 4097"), + Err(Error::TOO_LARGE) + ); + for line in [ + "vop search idx key k1 5", + "vop search idx vector 5 0", + "vop search idx vkey k1 10 filter", + "vop search idx vkey k1 10 withmeta withmeta", + ] { + assert_eq!(parse_line(line), Err(Error::MALFORMED), "{line}"); + } + } + + #[test] + fn keyed_commands_take_exact_arguments() { + assert_eq!(parse_line("vop getmeta idx k1").unwrap(), "getmeta idx k1"); + assert_eq!(parse_line("vop delete idx k1").unwrap(), "delete idx k1"); + assert_eq!(parse_line("vop drop idx").unwrap(), "drop idx"); + assert_eq!(parse_line("vop drop idx extra"), Err(Error::MALFORMED)); + assert_eq!(parse_line("vop"), Err(Error::MALFORMED)); + assert_eq!(parse_line("vop list"), Err(Error::MALFORMED)); + } +} diff --git a/src/protocol/tokens.rs b/src/protocol/tokens.rs new file mode 100644 index 0000000..d6bf5b2 --- /dev/null +++ b/src/protocol/tokens.rs @@ -0,0 +1,68 @@ +use std::os::raw::c_int; + +use super::Error; +use crate::memcached::ffi::token_t; + +#[derive(Clone, Copy)] +pub struct Tokens<'a> { + tokens: &'a [token_t], +} + +impl<'a> Tokens<'a> { + /// # Safety + /// + /// If `argc > 0`, `argv` must point to `argc` tokens with `length` readable bytes, for `'a`. + pub unsafe fn new(argv: *const token_t, argc: c_int) -> Self { + let len = usize::try_from(argc).unwrap_or(0); + if len == 0 || argv.is_null() { + return Self { tokens: &[] }; + } + // SAFETY: `len > 0` and `argv` is non-null, so it points to `len` tokens valid for `'a`. + let tokens = unsafe { std::slice::from_raw_parts(argv, len) }; + Self { tokens } + } + + pub fn len(&self) -> usize { + self.tokens.len() + } + + pub fn text(&self, i: usize) -> Result<&'a str, Error> { + let token = self.tokens.get(i).ok_or(Error::MALFORMED)?; + if token.value.is_null() || token.length == 0 { + return Ok(""); + } + // SAFETY: `value` is non-null, so it points to `length` bytes valid for `'a`. + let bytes = unsafe { std::slice::from_raw_parts(token.value.cast::(), token.length) }; + std::str::from_utf8(bytes).map_err(|_| Error::MALFORMED) + } +} + +/// The returned buffer owns the bytes the tokens point into, so it must outlive them. +#[cfg(test)] +pub fn tokenize(line: &str) -> (Vec, Vec) { + let mut buf = line.as_bytes().to_vec(); + let mut spans = Vec::new(); + let mut start = 0; + for (i, byte) in buf.iter_mut().enumerate() { + if *byte == b' ' { + if start != i { + spans.push((start, i - start)); + } + *byte = 0; + start = i + 1; + } + } + if start != buf.len() { + spans.push((start, buf.len() - start)); + } + let base = buf.as_mut_ptr(); + let tokens = spans + .into_iter() + .map(|(off, len)| token_t { + // SAFETY: every span starts inside `buf`, so `base.add(off)` stays in bounds. + value: unsafe { base.add(off) }.cast(), + length: len, + }) + .collect(); + (buf, tokens) +}