diff --git a/src/ascii.rs b/src/ascii.rs new file mode 100644 index 0000000..097cd93 --- /dev/null +++ b/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 = const { Cell::new(false) }; +} + +type ResponseHandler = + Option 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::(), 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 = 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 = 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`, 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::>(), + std::slice::from_raw_parts(dta.cast::(), nbytes as usize), + ) + }; + reply.extend_from_slice(bytes); + true + } + + #[test] + fn every_vop_line_is_not_supported() { + 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 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"); + } +} diff --git a/src/lib.rs b/src/lib.rs index 69aaf28..73d66a4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,63 @@ +//! 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(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)] @@ -9,8 +65,8 @@ 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); } }