From c80e8b862503f6dceb8eeeaacd72643368fbba50 Mon Sep 17 00:00:00 2001 From: ing-eoking Date: Thu, 29 Oct 2026 14:06:35 +0900 Subject: [PATCH] INTERNAL: Receive the body of a vop command --- src/ascii.rs | 238 ++++++++++++++++++++++++++++++++-------- src/cmd.rs | 30 +++++ src/index.rs | 4 + src/lib.rs | 2 + src/memcached/server.rs | 28 ++++- src/nread.rs | 155 ++++++++++++++++++++++++++ src/protocol.rs | 30 ++++- src/protocol/tokens.rs | 17 +++ 8 files changed, 457 insertions(+), 47 deletions(-) create mode 100644 src/cmd.rs create mode 100644 src/nread.rs diff --git a/src/ascii.rs b/src/ascii.rs index eb861ee..6a648a3 100644 --- a/src/ascii.rs +++ b/src/ascii.rs @@ -2,19 +2,26 @@ use std::cell::Cell; use std::os::raw::{c_char, c_int, c_void}; use std::ptr; -use crate::aio; -use crate::error::Error; +use crate::cmd; use crate::guard; -use crate::index; use crate::memcached::ffi::{AUTHZ_NONE, EXTENSION_ASCII_PROTOCOL_DESCRIPTOR, token_t}; use crate::memcached::server::{Connection, ResponseHandler}; -use crate::protocol::{self, Request, Tokens}; +use crate::nread::{self, Body}; +use crate::protocol::{self, Tokens}; const COMMAND: &[u8] = b"vop"; thread_local! { - /// Lets the `execute` the core calls right after a panicked `accept` close the connection. - static ACCEPT_PANICKED: Cell = const { Cell::new(false) }; + /// Set by [`close`] for the `execute` the core calls right after `accept`. + static CLOSE: Cell = const { Cell::new(false) }; +} + +/// Claims the line without its body and has the `execute` right after close the connection, the +/// only way to keep the stream in step when the body cannot be read. The core starts `ndata` at +/// 0, and `accept` sets it last, so it is still 0 here. +fn close() -> bool { + CLOSE.set(true); + true } pub(crate) static mut DESCRIPTOR: EXTENSION_ASCII_PROTOCOL_DESCRIPTOR = @@ -38,26 +45,19 @@ extern "C" fn get_auth_flag() -> u16 { /// # Safety /// -/// Non-null `argv` must be a token with `length` readable bytes; non-null `ndata` must be writable. +/// `argv`/`argc` must meet [`Tokens::new`]'s contract, `cookie` must be the live connection, and +/// non-null `ndata`/`ptr` must be writable. unsafe extern "C" fn accept( _cmd_cookie: *const c_void, - _cookie: *mut c_void, + cookie: *mut c_void, argc: c_int, argv: *mut token_t, ndata: *mut usize, - _ptr: *mut *mut c_char, + ptr: *mut *mut c_char, ) -> bool { // Clear a flag whose `execute` never ran (refused by SASL), or it closes another connection. - ACCEPT_PANICKED.set(false); - let claim_and_close = || { - if !ndata.is_null() { - // SAFETY: `ndata` is non-null, so the caller guarantees it is valid to write. - unsafe { *ndata = 0 }; - } - ACCEPT_PANICKED.set(true); - true - }; - guard(claim_and_close, || { + CLOSE.set(false); + guard(close, || { if argc < 1 || argv.is_null() { return false; } @@ -68,7 +68,31 @@ unsafe extern "C" fn accept( } // SAFETY: `value` is non-null, so it points to `length` readable bytes. let name = unsafe { std::slice::from_raw_parts(first.value.cast::(), first.length) }; - name == COMMAND + if name != COMMAND { + return false; + } + // SAFETY: the caller guarantees `Tokens::new`'s contract. + let tokens = unsafe { Tokens::new(argv, argc) }; + // `execute` answers a bad line. + let Ok(request) = protocol::parse(&tokens) else { + return true; + }; + let len = match nread::body(&request) { + Body::None => return true, + Body::Read(len) if !ndata.is_null() && !ptr.is_null() => len, + Body::Read(_) | Body::Refuse => return close(), + }; + // SAFETY: `cookie` is the live connection. + let Some(buffer) = (unsafe { nread::expect(Connection::new(cookie, None), &tokens, len) }) + else { + return close(); + }; + // SAFETY: both are non-null, so the caller guarantees they are writable. + unsafe { + *ndata = len; + *ptr = buffer; + } + true }) } @@ -85,34 +109,39 @@ unsafe extern "C" fn execute( guard( || false, || { - if ACCEPT_PANICKED.replace(false) { + if CLOSE.replace(false) { return false; } - // SAFETY: the caller guarantees `Tokens::new`'s contract. - let tokens = unsafe { Tokens::new(argv, argc) }; // SAFETY: the caller guarantees `response_handler` accepts `cookie`. let conn = unsafe { Connection::new(cookie, response_handler) }; - let reply = match protocol::parse(&tokens) { - Ok(Request::Search(_)) => { - let reply = Error::NotSupported.reply(); - // SAFETY: `conn` is running `execute`. - unsafe { aio::begin(conn, reply) }; - return true; - } - Ok(Request::Create(spec)) => match index::create(&spec) { - Ok(()) => b"CREATED\r\n", - Err(e) => e.reply(), - }, - Ok(_) => Error::NotSupported.reply(), - Err(e) => e.reply(), - }; + if argc == 0 { + // SAFETY: without a line, the core is done reading the body `accept` kept on `conn`. + let Some(pending) = (unsafe { nread::take(conn) }) else { + return false; + }; + // SAFETY: `conn` is running `execute`. + return unsafe { cmd::run(conn, pending.request()) }; + } + // SAFETY: the caller guarantees `Tokens::new`'s contract. + let tokens = unsafe { Tokens::new(argv, argc) }; // SAFETY: `conn` is running `execute`. - unsafe { conn.respond(reply) } + unsafe { cmd::run(conn, protocol::parse(&tokens)) } }, ) } -extern "C" fn abort(_cmd_cookie: *const c_void, _cookie: *const c_void) {} +/// # Safety +/// +/// `cookie` must be a live connection whose body the core stopped reading. +unsafe extern "C" fn abort(_cmd_cookie: *const c_void, cookie: *const c_void) { + guard( + || (), + || { + // SAFETY: the caller guarantees `cookie` is live, and only `accept` keeps on it. + drop(unsafe { nread::take(Connection::new(cookie, None)) }); + }, + ); +} #[cfg(test)] mod tests { @@ -142,7 +171,7 @@ mod tests { #[test] fn a_panicked_accept_closes_the_connection() { - ACCEPT_PANICKED.set(true); + CLOSE.set(true); let mut reply: Vec = Vec::new(); let mut tokens = vec![token("vop")]; // SAFETY: `record` gets the unshared `reply` back as its cookie. @@ -157,14 +186,14 @@ mod tests { }; assert!(!ok); assert!(reply.is_empty()); - assert!(!ACCEPT_PANICKED.get()); + assert!(!CLOSE.get()); } #[test] fn accept_clears_a_flag_left_behind() { - ACCEPT_PANICKED.set(true); + CLOSE.set(true); assert!(!accepts(&["get", "key"])); - assert!(!ACCEPT_PANICKED.get()); + assert!(!CLOSE.get()); } #[test] @@ -209,7 +238,8 @@ mod tests { #[test] fn well_formed_lines_are_not_supported_yet() { - assert_eq!(reply_to("vop drop idx"), "NOT_SUPPORTED\r\n"); + assert_eq!(reply_to("vop create unsupported 2 0"), "CREATED\r\n"); + assert_eq!(reply_to("vop drop unsupported"), "NOT_SUPPORTED\r\n"); } #[test] @@ -229,4 +259,124 @@ mod tests { "CLIENT_ERROR bad command line format\r\n" ); } + + /// Runs a body-carrying line through `accept`, the core's read, and `execute`. + fn reply_to_body(line: &str, body: &[u8]) -> Result { + crate::memcached::server::fake::install(); + let mut reply: Vec = Vec::new(); + let cookie: *mut c_void = (&raw mut reply).cast(); + let (_buf, mut tokens) = protocol::tokenize(line); + let (mut ndata, mut buffer) = (0, ptr::null_mut()); + // SAFETY: `tokens` point into `_buf`; `ndata`/`buffer` are writable. + let claimed = unsafe { + accept( + ptr::null(), + cookie, + tokens.len() as c_int, + tokens.as_mut_ptr(), + &raw mut ndata, + &raw mut buffer, + ) + }; + assert!(claimed); + if ndata == 0 { + // SAFETY: as above, and `record` gets `reply` back as its cookie. + let open = unsafe { + execute( + ptr::null(), + cookie, + tokens.len() as c_int, + tokens.as_mut_ptr(), + Some(record), + ) + }; + return if open { + Ok(String::from_utf8(reply).unwrap()) + } else { + Err("closed") + }; + } + assert_eq!(ndata, body.len()); + // SAFETY: `buffer` is the `ndata`-byte buffer `accept` handed out. + unsafe { ptr::copy_nonoverlapping(body.as_ptr(), buffer.cast(), body.len()) }; + // SAFETY: the body is in; `record` gets `reply` back as its cookie. + let open = unsafe { execute(ptr::null(), cookie, 0, ptr::null_mut(), Some(record)) }; + assert!(open); + Ok(String::from_utf8(reply).unwrap()) + } + + #[test] + fn a_body_is_read_then_answered() { + let _ = reply_to("vop create narrow 4 0"); + assert_eq!( + reply_to_body("vop setmeta narrow k1 5", b"hello\r\n"), + Ok("NOT_SUPPORTED\r\n".into()) + ); + assert_eq!( + reply_to_body("vop upsert narrow k1 4 meta 2", b"{}\r\nabcd\r\n"), + Ok("NOT_SUPPORTED\r\n".into()) + ); + } + + #[test] + fn a_segment_without_its_crlf_is_a_bad_data_chunk() { + assert_eq!( + reply_to_body("vop setmeta idx k1 5", b"hello!!"), + Ok("CLIENT_ERROR bad data chunk\r\n".into()) + ); + } + + #[test] + fn a_refused_body_closes_the_connection() { + let line = format!("vop setmeta idx k1 {}", protocol::MAX_BODY_BYTES); + assert_eq!(reply_to_body(&line, b""), Err("closed")); + } + + #[test] + fn a_missing_index_is_not_found() { + for line in [ + "vop drop missing", + "vop getmeta missing k1", + "vop delete missing k1", + "vop search missing vkey k1 5", + ] { + assert_eq!(reply_to(line), "NOT_FOUND\r\n", "{line}"); + } + assert_eq!( + reply_to_body("vop setmeta missing k1 5", b"hello\r\n"), + Ok("NOT_FOUND\r\n".into()) + ); + } + + #[test] + fn a_bad_line_is_answered_without_its_body() { + assert_eq!( + reply_to_body("vop setmeta idx k1", b""), + Ok("CLIENT_ERROR bad command line format\r\n".into()) + ); + } + + #[test] + fn abort_drops_the_kept_body() { + crate::memcached::server::fake::install(); + let mut reply: Vec = Vec::new(); + let cookie: *mut c_void = (&raw mut reply).cast(); + let (_buf, mut tokens) = protocol::tokenize("vop setmeta idx k1 5"); + let (mut ndata, mut buffer) = (0, ptr::null_mut()); + // SAFETY: `tokens` point into `_buf`; `ndata`/`buffer` are writable. + unsafe { + accept( + ptr::null(), + cookie, + tokens.len() as c_int, + tokens.as_mut_ptr(), + &raw mut ndata, + &raw mut buffer, + ); + abort(ptr::null(), cookie); + } + assert_eq!(ndata, 7); + // SAFETY: `cookie` stands for a live connection in the fake server. + assert!(unsafe { nread::take(Connection::new(cookie, None)) }.is_none()); + } } diff --git a/src/cmd.rs b/src/cmd.rs new file mode 100644 index 0000000..d6ce70a --- /dev/null +++ b/src/cmd.rs @@ -0,0 +1,30 @@ +use crate::aio; +use crate::error::Error; +use crate::index; +use crate::memcached::server::Connection; +use crate::protocol::Request; + +/// Runs `request` on `conn`, replying now or once a lane is done. +/// +/// # Safety +/// +/// `conn` must be running `execute`. +pub unsafe fn run(conn: Connection, request: Result, Error>) -> bool { + let reply = match request { + Ok(Request::Create(spec)) => match index::create(&spec) { + Ok(()) => b"CREATED\r\n", + Err(e) => e.reply(), + }, + Ok(request) if !index::exists(request.index()) => Error::NotFound.reply(), + Ok(Request::Search(_)) => { + let reply = Error::NotSupported.reply(); + // SAFETY: the caller guarantees `conn` is running `execute`. + unsafe { aio::begin(conn, reply) }; + return true; + } + Ok(_) => Error::NotSupported.reply(), + Err(e) => e.reply(), + }; + // SAFETY: the caller guarantees `conn` is running `execute`. + unsafe { conn.respond(reply) } +} diff --git a/src/index.rs b/src/index.rs index 7ed0334..d1db68a 100644 --- a/src/index.rs +++ b/src/index.rs @@ -23,6 +23,10 @@ pub fn create(spec: &Create<'_>) -> Result<(), Error> { Ok(()) } +pub fn exists(name: &str) -> bool { + read().contains_key(name) +} + fn read() -> RwLockReadGuard<'static, HashMap> { INDEXES.read().unwrap_or_else(PoisonError::into_inner) } diff --git a/src/lib.rs b/src/lib.rs index e4d8d5b..e493054 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,9 +2,11 @@ mod aio; mod ascii; +mod cmd; mod error; mod index; mod memcached; +mod nread; mod protocol; use std::os::raw::c_char; diff --git a/src/memcached/server.rs b/src/memcached/server.rs index 1e21893..b3b95fb 100644 --- a/src/memcached/server.rs +++ b/src/memcached/server.rs @@ -153,6 +153,22 @@ impl Connection { unsafe { handler(self.cookie, len, reply.as_ptr().cast()) } } + /// # Safety + /// + /// The connection must be live. + pub unsafe fn keep(self, data: *mut c_void) { + // SAFETY: the caller guarantees the connection is live. + unsafe { (api().store_engine_specific)(self.cookie, data) }; + } + + /// # Safety + /// + /// The connection must be live. + pub unsafe fn kept(self) -> *mut c_void { + // SAFETY: the caller guarantees the connection is live. + unsafe { (api().get_engine_specific)(self.cookie) } + } + /// Call before handing the work to another thread, so a notify that comes first still counts. /// /// # Safety @@ -191,6 +207,7 @@ impl Drop for IoWaiter { /// A core API for unit tests, where each cookie is a [`fake::Conn`]. #[cfg(test)] pub mod fake { + use std::cell::Cell; use std::mem; use std::ptr; use std::sync::mpsc::Sender; @@ -289,10 +306,17 @@ pub mod fake { 0 } - unsafe extern "C" fn store_engine_specific(_: *const c_void, _: *mut c_void) {} + thread_local! { + static KEPT: Cell<*mut c_void> = const { Cell::new(ptr::null_mut()) }; + } + + /// Keeps one value per thread, whatever the cookie. + unsafe extern "C" fn store_engine_specific(_: *const c_void, data: *mut c_void) { + KEPT.set(data); + } unsafe extern "C" fn get_engine_specific(_: *const c_void) -> *mut c_void { - ptr::null_mut() + KEPT.get() } unsafe extern "C" fn register_extension(_: extension_type_t, _: *mut c_void) -> bool { diff --git a/src/nread.rs b/src/nread.rs new file mode 100644 index 0000000..8892187 --- /dev/null +++ b/src/nread.rs @@ -0,0 +1,155 @@ +use std::alloc::{Layout, alloc}; +use std::os::raw::{c_char, c_int}; +use std::ptr; + +use crate::error::Error; +use crate::memcached::server::Connection; +use crate::protocol::{self, MAX_BODY_BYTES, Request, Tokens}; + +const MAX_TOKENS: usize = 16; + +pub enum Body { + None, + Read(usize), + /// Too long to read; the connection closes, as the core cannot skip it. + Refuse, +} + +/// What to do with the body `request` declares. +pub fn body(request: &Request<'_>) -> Body { + let segments = request.segments(); + if segments == [None, None] { + return Body::None; + } + let len = segments + .into_iter() + .flatten() + .try_fold(0usize, |sum, n| sum.checked_add(n)?.checked_add(2)); + match len { + Some(len) if len <= MAX_BODY_BYTES => Body::Read(len), + _ => Body::Refuse, + } +} + +/// A line kept on its connection while the core reads its body. +pub struct Pending { + line: Vec, + body: Vec, +} + +impl Pending { + /// Parses the kept line again and checks each body segment ends in CRLF. + pub fn request(&self) -> Result, Error> { + let (raw, n) = protocol::split::(&self.line).ok_or(Error::BadCommandLine)?; + // SAFETY: the tokens point into `self.line`, which outlives the borrow of `self`; `raw` dies + // here, but the request keeps only text from `self.line`, never the tokens. + let tokens = unsafe { Tokens::new(raw.as_ptr(), c_int::try_from(n).unwrap_or(0)) }; + let request = protocol::parse(&tokens)?; + let mut rest = self.body.as_slice(); + for n in request.segments().into_iter().flatten() { + let (segment, tail) = rest.split_at_checked(n + 2).ok_or(Error::BadDataChunk)?; + if !segment.ends_with(b"\r\n") { + return Err(Error::BadDataChunk); + } + rest = tail; + } + Ok(request) + } +} + +/// Keeps the line on `conn` with a `len`-byte body buffer, which it returns; `None` without memory, +/// or if `conn` already keeps something, which must not be lost. +/// +/// # Safety +/// +/// `conn` must be live. +pub unsafe fn expect(conn: Connection, tokens: &Tokens<'_>, len: usize) -> Option<*mut c_char> { + // SAFETY: the caller guarantees `conn` is live. + if !unsafe { conn.kept() }.is_null() { + return None; + } + let words = (0..tokens.len()).map(|i| tokens.text(i).ok()); + let mut line = Vec::new(); + line.try_reserve_exact(words.clone().map(|w| w.map_or(0, str::len) + 1).sum()) + .ok()?; + for word in words { + line.extend_from_slice(word?.as_bytes()); + line.push(b' '); + } + let mut body = Vec::new(); + body.try_reserve_exact(len).ok()?; + body.resize(len, 0); + let pending = try_box(Pending { line, body })?; + // SAFETY: `pending` is a live box; its body buffer stays put while the box is kept. + let buffer = unsafe { (*pending).body.as_mut_ptr() }.cast(); + // SAFETY: the caller guarantees `conn` is live. + unsafe { conn.keep(pending.cast()) }; + Some(buffer) +} + +/// # Safety +/// +/// `conn` must be live, and keep nothing but what [`expect`] kept. +pub unsafe fn take(conn: Connection) -> Option> { + // SAFETY: the caller guarantees `conn` is live. + let raw = unsafe { conn.kept() }.cast::(); + if raw.is_null() { + return None; + } + // SAFETY: as above. + unsafe { conn.keep(ptr::null_mut()) }; + // SAFETY: a non-null kept value is the box `expect` leaked, taken back once. + Some(unsafe { Box::from_raw(raw) }) +} + +/// Like `Box::into_raw(Box::new(..))`, but a failed allocation is `None` instead of an abort. +#[expect( + clippy::cast_ptr_alignment, + reason = "`alloc` aligns the memory for the layout it is given" +)] +fn try_box(pending: Pending) -> Option<*mut Pending> { + // SAFETY: `Pending` is not zero-sized. + let raw = unsafe { alloc(Layout::new::()) }.cast::(); + if raw.is_null() { + return None; + } + // SAFETY: `raw` is a fresh allocation with `Pending`'s layout. + unsafe { raw.write(pending) }; + Some(raw) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::tokenize; + + fn body_of(line: &str) -> Body { + let (_buf, raw) = tokenize(line); + // SAFETY: `raw` points into `_buf`, which outlives the call. + let tokens = unsafe { Tokens::new(raw.as_ptr(), c_int::try_from(raw.len()).unwrap()) }; + body(&protocol::parse(&tokens).unwrap()) + } + + #[test] + fn each_segment_counts_with_its_crlf() { + assert!(matches!(body_of("vop setmeta idx k1 5"), Body::Read(7))); + assert!(matches!( + body_of("vop upsert idx k1 8 meta 3"), + Body::Read(15) + )); + assert!(matches!( + body_of("vop search idx vector 5 8 filter 4"), + Body::Read(16) + )); + assert!(matches!(body_of("vop search idx vkey k1 5"), Body::None)); + assert!(matches!(body_of("vop drop idx"), Body::None)); + } + + #[test] + fn a_body_over_the_limit_is_refused() { + let line = format!("vop setmeta idx k1 {}", MAX_BODY_BYTES - 1); + assert!(matches!(body_of(&line), Body::Refuse)); + let line = format!("vop setmeta idx k1 {}", MAX_BODY_BYTES - 2); + assert!(matches!(body_of(&line), Body::Read(MAX_BODY_BYTES))); + } +} diff --git a/src/protocol.rs b/src/protocol.rs index 62033ca..82e6e86 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -15,9 +15,9 @@ use usearch::{MetricKind, ScalarKind}; use crate::error::Error; -pub use tokens::Tokens; #[cfg(test)] pub use tokens::tokenize; +pub use tokens::{Tokens, split}; pub const DEFAULT_M: usize = 16; pub const M_RANGE: RangeInclusive = 2..=1024; @@ -26,6 +26,7 @@ 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; +pub const MAX_BODY_BYTES: usize = 1 << 20; fn parse_metric(word: &str) -> Option { match word { @@ -136,6 +137,33 @@ pub enum Request<'a> { }, } +impl<'a> Request<'a> { + pub const fn index(&self) -> &'a str { + match self { + Self::Create(c) => c.index, + Self::Upsert(u) => u.index, + Self::SetMeta(s) => s.index, + Self::Search(s) => s.index, + Self::GetMeta { index, .. } | Self::Delete { index, .. } | Self::Drop { index } => { + index + } + } + } + + /// The lengths of the body segments the line declares, in the order they arrive. + pub const fn segments(&self) -> [Option; 2] { + match self { + Self::Upsert(u) => [u.meta_bytes, Some(u.vector_bytes)], + Self::SetMeta(s) => [Some(s.meta_bytes), None], + Self::Search(s) => match s.source { + Source::Vector { bytes, .. } => [s.filter_bytes, Some(bytes)], + Source::Vkey(_) => [s.filter_bytes, None], + }, + _ => [None, None], + } + } +} + fn in_range(n: T, range: &RangeInclusive) -> Result { if range.contains(&n) { Ok(n) diff --git a/src/protocol/tokens.rs b/src/protocol/tokens.rs index 960782e..fa64e48 100644 --- a/src/protocol/tokens.rs +++ b/src/protocol/tokens.rs @@ -37,6 +37,23 @@ impl<'a> Tokens<'a> { } } +/// Splits `line` on spaces into at most `N` tokens that point into it; `None` if there are more. +pub fn split(line: &[u8]) -> Option<([token_t; N], usize)> { + let mut tokens = [token_t { + value: std::ptr::null_mut(), + length: 0, + }; N]; + let mut n = 0; + for word in line.split(|&b| b == b' ').filter(|w| !w.is_empty()) { + *tokens.get_mut(n)? = token_t { + value: word.as_ptr().cast_mut().cast(), + length: word.len(), + }; + n += 1; + } + Some((tokens, n)) +} + /// The returned buffer owns the bytes the tokens point into, so it must outlive them. #[cfg(test)] pub fn tokenize(line: &str) -> (Vec, Vec) {