diff --git a/src/aio.rs b/src/aio.rs index f87a242..eac2246 100644 --- a/src/aio.rs +++ b/src/aio.rs @@ -34,8 +34,8 @@ unsafe extern "C" fn wake( || { // SAFETY: `data` is the job from `begin`; the core calls this once, after the lane. let job = unsafe { Box::from_raw(data.cast::()) }; - // SAFETY: the server passes the handler for `cookie`, which is in its io callback. - unsafe { Connection::new(cookie, handler).respond(job.reply) } + // SAFETY: the server is running the io callback for `cookie`, with `handler`. + unsafe { Connection::new(cookie, handler) }.respond(job.reply) }, ) } diff --git a/src/ascii.rs b/src/ascii.rs index eb861ee..1be2cd7 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"; +const COMMAND: &str = "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,37 +45,47 @@ 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, || { - if argc < 1 || argv.is_null() { + CLOSE.set(false); + guard(close, || { + // SAFETY: the caller guarantees `Tokens::new`'s contract. + let Ok(tokens) = (unsafe { Tokens::new(argv, argc) }) else { return false; - } - // SAFETY: `argc >= 1` and `argv` is non-null, so it points to a valid token. - let first = unsafe { &*argv }; - if first.value.is_null() { + }; + if tokens.words().first() != Some(&COMMAND) { return false; } - // 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 + // `execute` answers a bad line. + let Ok(request) = protocol::parse(tokens.words()) else { + return true; + }; + let len = match nread::body(&request) { + Body::None => return true, + Body::Read(len) => len, + Body::Refuse => return close(), + }; + // SAFETY: the caller guarantees non-null ones are writable. + let Some((ndata, ptr)) = (unsafe { ndata.as_mut().zip(ptr.as_mut()) }) else { + return close(); + }; + // SAFETY: the server is running `accept` for `cookie`. + let conn = unsafe { Connection::new(cookie, None) }; + let Some(buffer) = nread::expect(conn, tokens.words(), len) else { + return close(); + }; + *ndata = len; + *ptr = buffer; + true }) } @@ -85,38 +102,44 @@ 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`. + // SAFETY: the server is running `execute` for `cookie`, with `response_handler`. 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 { + let Some(pending) = conn.take_kept() else { + return false; + }; + // SAFETY: `conn` is running `execute`. + return unsafe { cmd::run(conn, pending.request()) }; + } + // SAFETY: the caller guarantees `Tokens::new`'s contract. + let request = unsafe { Tokens::new(argv, argc) } + .and_then(|tokens| protocol::parse(tokens.words())); // SAFETY: `conn` is running `execute`. - unsafe { conn.respond(reply) } + unsafe { cmd::run(conn, request) } }, ) } -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 server is running `abort` for `cookie`. + drop(unsafe { Connection::new(cookie, None) }.take_kept()); + }, + ); +} #[cfg(test)] mod tests { use super::*; + use crate::protocol::tokenize; fn token(text: &str) -> token_t { token_t { @@ -142,7 +165,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 +180,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] @@ -192,7 +215,7 @@ mod tests { fn reply_to(line: &str) -> String { let mut reply: Vec = Vec::new(); - let (_buf, mut tokens) = protocol::tokenize(line); + let (_buf, mut tokens) = tokenize(line); // SAFETY: `tokens` point into `_buf`, and `record` gets the unshared `reply` back. let ok = unsafe { execute( @@ -209,7 +232,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 +253,125 @@ 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) = 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) = 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. + let conn = unsafe { Connection::new(cookie, None) }; + assert!(conn.take_kept().is_none()); + } } diff --git a/src/cmd.rs b/src/cmd.rs new file mode 100644 index 0000000..cc0b260 --- /dev/null +++ b/src/cmd.rs @@ -0,0 +1,29 @@ +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(), + }; + conn.respond(reply) +} diff --git a/src/index.rs b/src/index.rs index 7ed0334..77cfa3e 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) } @@ -45,16 +49,11 @@ fn options(spec: &Create<'_>) -> IndexOptions { #[cfg(test)] mod tests { - use std::os::raw::c_int; - use super::*; - use crate::protocol::{self, Request}; + use crate::protocol::{self, Request, Tokens}; fn create_line(line: &str) -> Result<(), Error> { - let (_buf, tokens) = protocol::tokenize(line); - // SAFETY: `tokens` point into `_buf`, which outlives the call. - let tokens = unsafe { protocol::Tokens::new(tokens.as_ptr(), tokens.len() as c_int) }; - let Ok(Request::Create(spec)) = protocol::parse(&tokens) else { + let Ok(Request::Create(spec)) = protocol::parse(Tokens::from_line(line).words()) else { panic!("not a create line: {line}"); }; create(&spec) 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..55664b1 100644 --- a/src/memcached/server.rs +++ b/src/memcached/server.rs @@ -1,5 +1,6 @@ use std::ffi::CStr; use std::os::raw::{c_char, c_int, c_void}; +use std::ptr; use std::sync::OnceLock; use super::ffi::{ @@ -8,6 +9,7 @@ use super::ffi::{ EXTENSION_LOGGER_DESCRIPTOR, SERVER_HANDLE_V1, VERSION, extension_type_t, extension_type_t_EXTENSION_ASCII_PROTOCOL, rel_time_t, time_t, }; +use crate::nread::Pending; pub type ResponseHandler = Option bool>; @@ -133,26 +135,55 @@ pub struct Connection { impl Connection { /// # Safety /// - /// `handler`, if present, must accept `cookie`. + /// The server must be running `accept`, `execute`, `abort` or the io callback for `cookie`, + /// with the `handler` it passed (`None` for `accept` and `abort`), and the connection must not + /// outlive that call. pub const unsafe fn new(cookie: *const c_void, handler: ResponseHandler) -> Self { Self { cookie, handler } } - /// # Safety - /// - /// The connection must be running `execute` or its io callback. - pub unsafe fn respond(self, reply: &[u8]) -> bool { + pub fn respond(self, reply: &[u8]) -> bool { let Some(handler) = self.handler else { return false; }; let Ok(len) = c_int::try_from(reply.len()) else { return false; }; - // SAFETY: `new`'s caller guaranteed `handler` accepts the connection, which the caller - // guarantees can take a reply now; `reply` outlives the call. + // SAFETY: `new`'s caller guaranteed the connection takes a reply through `handler` now; + // `reply` outlives the call. unsafe { handler(self.cookie, len, reply.as_ptr().cast()) } } + /// Keeps `pending` on the connection until [`Connection::take_kept`]; gives it back if the + /// connection already keeps one. + pub fn keep(self, pending: Pending) -> Result<(), Pending> { + if !self.kept().is_null() { + return Err(pending); + } + self.store(Box::into_raw(Box::new(pending)).cast()); + Ok(()) + } + + pub fn take_kept(self) -> Option { + let kept = self.kept().cast::(); + if kept.is_null() { + return None; + } + self.store(ptr::null_mut()); + // SAFETY: only `keep` stores on a connection, a leaked box, taken back once here. + Some(*unsafe { Box::from_raw(kept) }) + } + + fn store(self, data: *mut c_void) { + // SAFETY: `new`'s caller guaranteed the connection is live. + unsafe { (api().store_engine_specific)(self.cookie, data) }; + } + + fn kept(self) -> *mut c_void { + // SAFETY: `new`'s caller guaranteed 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 +222,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; @@ -220,7 +252,7 @@ pub mod fake { } pub fn connection(&self) -> Connection { - // SAFETY: `record` takes any `Conn` cookie. + // SAFETY: `record` takes any `Conn` cookie, which takes a reply or a wait at any time. unsafe { Connection::new(ptr::from_ref(self).cast(), Some(record)) } } @@ -289,10 +321,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..1b882ad --- /dev/null +++ b/src/nread.rs @@ -0,0 +1,101 @@ +use std::os::raw::c_char; + +use crate::error::Error; +use crate::memcached::server::Connection; +use crate::protocol::{self, MAX_BODY_BYTES, Request, Tokens}; + +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: String, + body: Vec, +} + +impl Pending { + /// Parses the kept line again and checks each body segment ends in CRLF. + pub fn request(&self) -> Result, Error> { + let request = protocol::parse(Tokens::from_line(&self.line).words())?; + 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. +pub fn expect(conn: Connection, words: &[&str], len: usize) -> Option<*mut c_char> { + let mut line = String::new(); + line.try_reserve_exact(words.iter().map(|w| w.len() + 1).sum()) + .ok()?; + for word in words { + line.push_str(word); + line.push(' '); + } + let mut body = Vec::new(); + body.try_reserve_exact(len).ok()?; + body.resize(len, 0); + // The buffer stays put as `body` moves onto `conn`, until it is taken back and dropped. + let buffer = body.as_mut_ptr().cast(); + conn.keep(Pending { line, body }).ok()?; + Some(buffer) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn body_of(line: &str) -> Body { + body(&protocol::parse(Tokens::from_line(line).words()).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..be976bf 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -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) @@ -144,37 +172,41 @@ fn in_range(n: T, range: &RangeInclusive) -> Result } } -/// Reads the tokens of a line in order. -struct Args<'t, 'a> { - tokens: &'t Tokens<'a>, +/// Reads the words of a line in order. +struct Args<'w, 'a> { + words: &'w [&'a str], at: usize, } impl<'a> Args<'_, 'a> { fn is_empty(&self) -> bool { - self.at >= self.tokens.len() + self.at >= self.words.len() + } + + fn peek(&self) -> Option<&'a str> { + self.words.get(self.at).copied() } fn next(&mut self) -> Result<&'a str, Error> { - let text = self.tokens.text(self.at)?; + let word = self.peek().ok_or(Error::BadCommandLine)?; self.at += 1; - Ok(text) + Ok(word) } fn number(&mut self) -> Result { self.next()?.parse().map_err(|_| Error::BadCommandLine) } - /// Reads the next token only if it is a number. + /// Reads the next word only if it is a number. fn optional_number(&mut self) -> Option { - let n = self.tokens.text(self.at).ok()?.parse().ok()?; + let n = self.peek()?.parse().ok()?; self.at += 1; Some(n) } } -pub fn parse<'a>(tokens: &Tokens<'a>) -> Result, Error> { - let mut args = Args { tokens, at: 1 }; +pub fn parse<'a>(words: &[&'a str]) -> Result, Error> { + let mut args = Args { words, at: 1 }; let request = match args.next()? { "create" => Request::Create(create(&mut args)?), "upsert" => Request::Upsert(upsert(&mut args)?), @@ -319,15 +351,10 @@ fn search<'a>(args: &mut Args<'_, 'a>) -> Result, Error> { #[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)) + parse(Tokens::from_line(line).words()).map(|request| describe(&request)) } fn describe(request: &Request<'_>) -> String { diff --git a/src/protocol/tokens.rs b/src/protocol/tokens.rs index 960782e..6dd1661 100644 --- a/src/protocol/tokens.rs +++ b/src/protocol/tokens.rs @@ -3,37 +3,47 @@ use std::os::raw::c_int; use crate::error::Error; use crate::memcached::ffi::token_t; -#[derive(Clone, Copy)] +/// The words of a command line, in place. pub struct Tokens<'a> { - tokens: &'a [token_t], + words: Vec<&'a str>, } impl<'a> Tokens<'a> { + /// The tokens the core split a line into. A bad command line if one is not UTF-8. + /// /// # 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 { + pub unsafe fn new(argv: *const token_t, argc: c_int) -> Result { let len = usize::try_from(argc).unwrap_or(0); if len == 0 || argv.is_null() { - return Self { tokens: &[] }; + return Ok(Self { words: Vec::new() }); } // 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 } + let words = tokens + .iter() + .map(|token| { + if token.value.is_null() || token.length == 0 { + return Ok(""); + } + // SAFETY: the caller guarantees a non-null `value` has `length` bytes for `'a`. + let bytes = + unsafe { std::slice::from_raw_parts(token.value.cast::(), token.length) }; + str::from_utf8(bytes).map_err(|_| Error::BadCommandLine) + }) + .collect::>()?; + Ok(Self { words }) } - pub fn len(&self) -> usize { - self.tokens.len() + /// Splits `line` on spaces. + pub fn from_line(line: &'a str) -> Self { + let words = line.split(' ').filter(|w| !w.is_empty()).collect(); + Self { words } } - pub fn text(&self, i: usize) -> Result<&'a str, Error> { - let token = self.tokens.get(i).ok_or(Error::BadCommandLine)?; - 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::BadCommandLine) + pub fn words(&self) -> &[&'a str] { + &self.words } } @@ -66,3 +76,28 @@ pub fn tokenize(line: &str) -> (Vec, Vec) { .collect(); (buf, tokens) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_token_that_is_not_utf8_makes_a_bad_line() { + let token = |bytes: &[u8]| token_t { + value: bytes.as_ptr().cast_mut().cast(), + length: bytes.len(), + }; + let raw = [token(b"vop"), token(b"drop"), token(&[0xff])]; + // SAFETY: `raw` points into byte literals, which live for the whole program. + let tokens = unsafe { Tokens::new(raw.as_ptr(), raw.len() as c_int) }; + assert!(matches!(tokens, Err(Error::BadCommandLine))); + } + + #[test] + fn a_line_splits_on_any_run_of_spaces() { + assert_eq!( + Tokens::from_line(" vop drop idx ").words(), + ["vop", "drop", "idx"] + ); + } +}