INTERNAL: Receive the body of a vop command

Cargo.toml

@@ -53,5 +53,8 @@ unimplemented = "warn"
 use_self = "warn"
 or_fun_call = "warn"
 
+[dependencies]
+usearch = "2.26"
+
 [build-dependencies]
 bindgen = "0.73.2"

src/ascii.rs

@@ -2,18 +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::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<bool> = const { Cell::new(false) };
+    /// Set by [`close`] for the `execute` the core calls right after `accept`.
+    static CLOSE: Cell<bool> = 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 =
@@ -37,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;
         }
@@ -67,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::<u8>(), 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, which keeps nothing else.
+        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
     })
 }
 
@@ -84,30 +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(_) => 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 {
@@ -137,7 +171,7 @@ mod tests {
 
     #[test]
     fn a_panicked_accept_closes_the_connection() {
-        ACCEPT_PANICKED.set(true);
+        CLOSE.set(true);
         let mut reply: Vec<u8> = Vec::new();
         let mut tokens = vec![token("vop")];
         // SAFETY: `record` gets the unshared `reply` back as its cookie.
@@ -152,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]
@@ -204,8 +238,14 @@ mod tests {
 
     #[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");
+        assert_eq!(reply_to("vop create unsupported 2 0"), "CREATED\r\n");
+        assert_eq!(reply_to("vop drop unsupported"), "NOT_SUPPORTED\r\n");
+    }
+
+    #[test]
+    fn create_registers_the_index_once() {
+        assert_eq!(reply_to("vop create ascii 2 0"), "CREATED\r\n");
+        assert_eq!(reply_to("vop create ascii 2 0"), "EXISTS\r\n");
     }
 
     #[test]
@@ -219,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<String, &'static str> {
+        crate::memcached::server::fake::install();
+        let mut reply: Vec<u8> = 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<u8> = 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());
+    }
 }

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<Request<'_>, 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) }
+}