INTERNAL: Register the vop ASCII extension

src/ascii.rs

@@ -0,0 +1,200 @@
+use std::cell::Cell;
+use std::os::raw::{c_char, c_int, c_void};
+use std::ptr;
+
+use crate::guard;
+use crate::memcached::ffi::{AUTHZ_NONE, EXTENSION_ASCII_PROTOCOL_DESCRIPTOR, token_t};
+
+const COMMAND: &[u8] = b"vop";
+const NOT_SUPPORTED: &[u8] = b"NOT_SUPPORTED\r\n";
+
+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) };
+}
+
+type ResponseHandler =
+    Option<unsafe extern "C" fn(cookie: *const c_void, nbytes: c_int, dta: *const c_char) -> bool>;
+
+pub(crate) static mut DESCRIPTOR: EXTENSION_ASCII_PROTOCOL_DESCRIPTOR =
+    EXTENSION_ASCII_PROTOCOL_DESCRIPTOR {
+        get_name: Some(get_name),
+        get_auth_flag: Some(get_auth_flag),
+        accept: Some(accept),
+        execute: Some(execute),
+        abort: Some(abort),
+        cookie: ptr::null(),
+        next: ptr::null_mut(),
+    };
+
+extern "C" fn get_name(_cmd_cookie: *const c_void) -> *const c_char {
+    c"arcus_vector".as_ptr()
+}
+
+extern "C" fn get_auth_flag() -> u16 {
+    AUTHZ_NONE as u16
+}
+
+/// # Safety
+///
+/// Non-null `argv` must be a token with `length` readable bytes; non-null `ndata` must be writable.
+unsafe extern "C" fn accept(
+    _cmd_cookie: *const c_void,
+    _cookie: *mut c_void,
+    argc: c_int,
+    argv: *mut token_t,
+    ndata: *mut usize,
+    _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() {
+            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() {
+            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::<u8>(), first.length) };
+        name == COMMAND
+    })
+}
+
+/// # Safety
+///
+/// `response_handler`, if present, must accept `cookie` and bytes that live for the call.
+unsafe extern "C" fn execute(
+    _cmd_cookie: *const c_void,
+    cookie: *const c_void,
+    _argc: c_int,
+    _argv: *mut token_t,
+    response_handler: ResponseHandler,
+) -> bool {
+    guard(
+        || false,
+        || {
+            if ACCEPT_PANICKED.replace(false) {
+                return false;
+            }
+            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(),
+                )
+            }
+        },
+    )
+}
+
+extern "C" fn abort(_cmd_cookie: *const c_void, _cookie: *const c_void) {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn token(text: &str) -> token_t {
+        token_t {
+            value: text.as_ptr().cast_mut().cast(),
+            length: text.len(),
+        }
+    }
+
+    fn accepts(line: &[&str]) -> bool {
+        let mut tokens: Vec<token_t> = line.iter().map(|t| token(t)).collect();
+        // SAFETY: `tokens` point into `line`, which outlives the call.
+        unsafe {
+            accept(
+                ptr::null(),
+                ptr::null_mut(),
+                line.len() as c_int,
+                tokens.as_mut_ptr(),
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        }
+    }
+
+    #[test]
+    fn a_panicked_accept_closes_the_connection() {
+        ACCEPT_PANICKED.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.
+        let ok = unsafe {
+            execute(
+                ptr::null(),
+                (&raw mut reply).cast(),
+                1,
+                tokens.as_mut_ptr(),
+                Some(record),
+            )
+        };
+        assert!(!ok);
+        assert!(reply.is_empty());
+        assert!(!ACCEPT_PANICKED.get());
+    }
+
+    #[test]
+    fn accept_clears_a_flag_left_behind() {
+        ACCEPT_PANICKED.set(true);
+        assert!(!accepts(&["get", "key"]));
+        assert!(!ACCEPT_PANICKED.get());
+    }
+
+    #[test]
+    fn accepts_only_vop_lines() {
+        assert!(accepts(&["vop", "create", "idx", "2", "0"]));
+        assert!(!accepts(&["get", "key"]));
+        assert!(!accepts(&["vopx"]));
+        assert!(!accepts(&[]));
+    }
+
+    /// # Safety
+    ///
+    /// `cookie` must be an unshared `Vec<u8>`, and `dta` must hold `nbytes` bytes.
+    unsafe extern "C" fn record(cookie: *const c_void, nbytes: c_int, dta: *const c_char) -> bool {
+        // SAFETY: the caller guarantees both, per `# Safety`.
+        let (reply, bytes) = unsafe {
+            (
+                &mut *cookie.cast_mut().cast::<Vec<u8>>(),
+                std::slice::from_raw_parts(dta.cast::<u8>(), nbytes as usize),
+            )
+        };
+        reply.extend_from_slice(bytes);
+        true
+    }
+
+    #[test]
+    fn every_vop_line_is_not_supported() {
+        let mut reply: Vec<u8> = Vec::new();
+        let mut tokens = vec![token("vop"), token("create")];
+        // SAFETY: `record` gets the unshared `reply` back as its cookie.
+        let ok = unsafe {
+            execute(
+                ptr::null(),
+                (&raw mut reply).cast(),
+                2,
+                tokens.as_mut_ptr(),
+                Some(record),
+            )
+        };
+        assert!(ok);
+        assert_eq!(reply, b"NOT_SUPPORTED\r\n");
+    }
+}

src/lib.rs

@@ -1,16 +1,72 @@
+//! arcus-vector: vector search as an extension to arcus-memcached.
+
+mod ascii;
 mod memcached;
 
-pub fn add(left: u64, right: u64) -> u64 {
-    left + right
+use std::os::raw::c_char;
+use std::panic::{AssertUnwindSafe, catch_unwind};
+
+use crate::memcached::ffi::{
+    EXTENSION_ERROR_CODE, EXTENSION_ERROR_CODE_EXTENSION_FATAL,
+    EXTENSION_ERROR_CODE_EXTENSION_SUCCESS, GET_SERVER_API,
+    extension_type_t_EXTENSION_ASCII_PROTOCOL,
+};
+
+pub(crate) fn guard<T>(on_panic: impl FnOnce() -> T, f: impl FnOnce() -> T) -> T {
+    catch_unwind(AssertUnwindSafe(f)).unwrap_or_else(|_| on_panic())
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn memcached_extensions_initialize(
+    _config: *const c_char,
+    get_server_api: GET_SERVER_API,
+) -> EXTENSION_ERROR_CODE {
+    guard(
+        || EXTENSION_ERROR_CODE_EXTENSION_FATAL,
+        || initialize(get_server_api),
+    )
+}
+
+fn initialize(get_server_api: GET_SERVER_API) -> EXTENSION_ERROR_CODE {
+    let Some(get_server_api) = get_server_api else {
+        return EXTENSION_ERROR_CODE_EXTENSION_FATAL;
+    };
+    // SAFETY: the core's own function; it takes no arguments and returns its static handle.
+    let server = unsafe { get_server_api() };
+    if server.is_null() {
+        return EXTENSION_ERROR_CODE_EXTENSION_FATAL;
+    }
+    // SAFETY: `server` was null-checked above and points to the core's static handle.
+    let extension = unsafe { (*server).extension };
+    if extension.is_null() {
+        return EXTENSION_ERROR_CODE_EXTENSION_FATAL;
+    }
+    // SAFETY: `extension` was null-checked above and points to the core's static extension table.
+    let Some(register) = (unsafe { (*extension).register_extension }) else {
+        return EXTENSION_ERROR_CODE_EXTENSION_FATAL;
+    };
+
+    // SAFETY: `DESCRIPTOR` is static and only the core touches it from here on (it writes `next`).
+    let registered = unsafe {
+        register(
+            extension_type_t_EXTENSION_ASCII_PROTOCOL,
+            (&raw mut ascii::DESCRIPTOR).cast(),
+        )
+    };
+    if registered {
+        EXTENSION_ERROR_CODE_EXTENSION_SUCCESS
+    } else {
+        EXTENSION_ERROR_CODE_EXTENSION_FATAL
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn it_works() {
-        let result = add(2, 2);
-        assert_eq!(result, 4);
+    fn guard_turns_a_panic_into_the_fallback() {
+        assert_eq!(guard(|| 0, || 1), 1);
+        assert_eq!(guard(|| 0, || panic!("in a callback")), 0);
     }
 }