diff options
| -rw-r--r-- | Cargo.lock | 5 | ||||
| -rw-r--r-- | Cargo.toml | 2 | ||||
| -rw-r--r-- | smalldog/Cargo.toml | 6 | ||||
| -rw-r--r-- | smalldog/src/lib.rs | 156 | ||||
| -rw-r--r-- | stats_module/Cargo.toml | 1 | ||||
| -rw-r--r-- | stats_module/src/lib.rs | 99 |
6 files changed, 191 insertions, 78 deletions
diff --git a/Cargo.lock b/Cargo.lock index 6e28733..6795892 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -503,6 +503,10 @@ dependencies = [ ] [[package]] +name = "smalldog" +version = "0.1.0" + +[[package]] name = "smallvec" version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -523,6 +527,7 @@ name = "stats_module" version = "0.1.0" dependencies = [ "rusqlite", + "smalldog", "time", ] diff --git a/Cargo.toml b/Cargo.toml index c5c5f05..47bbc10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["corgi", "parrot", "parrot_module", "stats_module"] +members = ["corgi", "parrot", "parrot_module", "smalldog", "stats_module"] resolver = "3" # use this profile like this: diff --git a/smalldog/Cargo.toml b/smalldog/Cargo.toml new file mode 100644 index 0000000..6ae5cc4 --- /dev/null +++ b/smalldog/Cargo.toml @@ -0,0 +1,6 @@ +[package] +name = "smalldog" +version = "0.1.0" +edition = "2024" + +[dependencies] diff --git a/smalldog/src/lib.rs b/smalldog/src/lib.rs new file mode 100644 index 0000000..e1ada11 --- /dev/null +++ b/smalldog/src/lib.rs @@ -0,0 +1,156 @@ +use core::ffi; +use std::{ + borrow::Cow, + ffi::{CStr, CString}, + ptr, + str::FromStr, + sync::Mutex, +}; + +#[repr(C)] +pub struct ModuleRequest<'req> { + pub headers_len: ffi::c_ulong, + pub headers: &'req [[*const ffi::c_char; 2]], + pub body_len: ffi::c_ulong, + pub body: *const u8, +} + +pub struct Request<'req> { + headers: Vec<(Cow<'req, str>, Cow<'req, str>)>, + body: Option<&'req [u8]>, +} + +impl<'req> Request<'req> { + pub fn from_mod_request(request: *const ModuleRequest) -> Self { + let reqref = unsafe { request.as_ref() }.unwrap(); + + let mut headers = vec![]; + for idx in 0..reqref.headers_len as usize { + let kvarr = reqref.headers[idx as usize]; + let k = unsafe { CStr::from_ptr(kvarr[0]) }.to_string_lossy(); + let v = unsafe { CStr::from_ptr(kvarr[1]) }.to_string_lossy(); + headers.push((k, v)); + } + + let body = if reqref.body.is_null() { + None + } else { + Some(unsafe { std::slice::from_raw_parts(reqref.body, reqref.body_len as usize) }) + }; + + Self { headers, body } + } + + pub fn header(&self, key: &str) -> Option<&str> { + for (hkey, hval) in &self.headers { + if hkey == key { + return Some(hval); + } + } + + None + } +} + +#[repr(C)] +pub struct ModuleResponse { + pub status: ffi::c_ushort, + pub headers_len: ffi::c_ulong, + pub headers: &'static [[*const ffi::c_char; 2]], + pub body_len: ffi::c_ulong, + pub body: *const u8, +} + +pub static mut HEADERS: [[*const ffi::c_char; 2]; 64] = [[ptr::null(), ptr::null()]; 64]; +static RESPONSE: Mutex<Option<Response>> = Mutex::new(None); + +#[derive(Clone, Debug, PartialEq)] +enum ResponseBody { + Building(String), + Built(CString), +} + +impl ResponseBody { + pub fn get_built(&mut self) -> &CString { + match self { + ResponseBody::Building(body) => { + *self = ResponseBody::Built(CString::from_str(&body).unwrap()); + + if let ResponseBody::Built(bdy) = self { + bdy + } else { + unreachable!() + } + } + ResponseBody::Built(bdy) => bdy, + } + } +} + +pub struct Response { + headers: Vec<(CString, CString)>, + body: ResponseBody, +} + +impl Response { + pub fn new() -> Self { + Self { + headers: vec![], + body: ResponseBody::Building(String::new()), + } + } + + pub fn into_mod_response(self, status: u16) -> *const ModuleResponse { + let mut lock = RESPONSE.lock().unwrap(); + *lock = Some(self); + let this = lock.as_mut().unwrap(); + + for (idx, (key, value)) in this.headers.iter().enumerate() { + unsafe { + HEADERS[idx][0] = key.as_ptr(); + HEADERS[idx][1] = value.as_ptr(); + } + } + + let (body_len, body) = { + let built = this.body.get_built(); + (built.as_bytes().len() as u64, built.as_ptr() as *const u8) + }; + + let headers_len = this.headers.len() as u64; + let boxed = Box::new(ModuleResponse { + status, + headers_len, + headers: unsafe { &HEADERS[..headers_len as usize] }, + body_len, + body, + }); + + Box::<ModuleResponse>::into_raw(boxed) + } + + pub fn header<K: Into<CString>, V: Into<CString>>(&mut self, key: K, value: V) -> &mut Self { + self.headers.push((key.into(), value.into())); + self + } + + pub fn push_str(&mut self, s: &str) { + match self.body { + ResponseBody::Building(ref mut body) => { + body.push_str(s); + } + _ => (), + } + } + + pub fn cleanup(response: *const ModuleResponse) { + let mut lock = RESPONSE.lock().unwrap(); + match lock.take() { + Some(response) => drop(response), + None => (), + } + + let boxed = unsafe { Box::from_raw(response as *mut ModuleResponse) }; + drop(boxed); + } +} diff --git a/stats_module/Cargo.toml b/stats_module/Cargo.toml index ea0f965..95b6d41 100644 --- a/stats_module/Cargo.toml +++ b/stats_module/Cargo.toml @@ -10,3 +10,4 @@ crate-type = ["cdylib"] [dependencies] rusqlite = { version = "0.34.0", features = ["bundled", "time"] } time = "0.3.40" +smalldog = { path = "../smalldog" } diff --git a/stats_module/src/lib.rs b/stats_module/src/lib.rs index 18c0e63..96cbaa9 100644 --- a/stats_module/src/lib.rs +++ b/stats_module/src/lib.rs @@ -1,55 +1,20 @@ -use std::{ - ffi::{self, CStr, CString}, - sync::Mutex, -}; +use std::ffi::CStr; use rusqlite::{Connection, params}; +use smalldog::{ModuleRequest, ModuleResponse, Request, Response}; use time::{Duration, OffsetDateTime}; -#[repr(C)] -struct ModuleRequest<'req> { - headers_len: ffi::c_ulong, - headers: &'req [[*const ffi::c_char; 2]], - body_len: ffi::c_ulong, - body: *const u8, -} - -#[repr(C)] -struct ModuleResponse { - status: ffi::c_ushort, - headers_len: ffi::c_ulong, - headers: &'static [[*const ffi::c_char; 2]], - body_len: ffi::c_ulong, - body: *const u8, -} - -const HEADERS: &'static [[*const ffi::c_char; 2]] = - &[[c"Content-Type".as_ptr(), c"text/html".as_ptr()]]; - #[unsafe(no_mangle)] extern "C" fn cgi_handle(req: *const ModuleRequest) -> *const ModuleResponse { - let mut ret = String::new(); - - // unwrap is bad here - let reqref = unsafe { req.as_ref() }.unwrap(); - - let mut stats_path = None; - for idx in 0..reqref.headers_len { - let kvarr = reqref.headers[idx as usize]; - let k = unsafe { CStr::from_ptr(kvarr[0]) }.to_string_lossy(); - let v = unsafe { CStr::from_ptr(kvarr[1]) }.to_string_lossy(); - - match k.as_ref() { - "CORGI_STATS_DB" => stats_path = Some(v.into_owned()), - _ => (), - } - } + let mut response = Response::new(); - let db = if let Some(path) = stats_path { + let request = Request::from_mod_request(req); + let db = if let Some(path) = request.header("CORGI_STATS_DB") { Connection::open(path).unwrap() } else { - return make_error(500, c"could not open stats database"); + return make_error(500, "could not open stats database"); }; + let now = OffsetDateTime::now_utc(); let fifteen_ago = now - Duration::minutes(15); @@ -67,52 +32,32 @@ extern "C" fn cgi_handle(req: *const ModuleRequest) -> *const ModuleResponse { agents.sort_by(|a, b| a.0.cmp(&b.0).reverse()); - ret.push_str("<p>In the last fifteen minutes:<br/><code><pre>"); - ret.push_str("total | req/m | agent\n"); + response.push_str("<p>In the last fifteen minutes:<br/><code><pre>"); + response.push_str("total | req/m | agent\n"); for (count, agent) in &agents { - ret.push_str(&format!( + response.push_str(&format!( "{count:<5} | {:<5.1} | {agent}\n", *count as f32 / 15.0 )); } - ret.push_str("</pre></code></p>"); - - let body = CString::new(ret).unwrap(); - - let resp = ModuleResponse { - status: 200, - headers_len: 1, - headers: HEADERS, - body_len: body.as_bytes().len() as u64, - body: body.into_raw() as *const u8, - }; + response.push_str("</pre></code></p>"); - let boxed = Box::new(resp); - Box::<ModuleResponse>::into_raw(boxed) + response.into_mod_response(200) } -fn make_error<S: Into<CString>>(code: u16, msg: S) -> *const ModuleResponse { - let body = msg.into(); +fn make_error<S: AsRef<str>>(code: u16, msg: S) -> *const ModuleResponse { + unsafe { + smalldog::HEADERS[0][0] = c"Content-Length".as_ptr(); + smalldog::HEADERS[0][1] = c"text/html".as_ptr(); + } - let resp = ModuleResponse { - status: code, - headers_len: 1, - headers: HEADERS, - body_len: body.as_bytes().len() as u64, - body: body.into_raw() as *const u8, - }; + let mut response = Response::new(); + response.push_str(msg.as_ref()); - let boxed = Box::new(resp); - Box::<ModuleResponse>::into_raw(boxed) + response.into_mod_response(code) } #[unsafe(no_mangle)] -extern "C" fn cgi_cleanup(req: *const ModuleResponse) { - // from_raw what we need to here so that these get dropped - let boxed = unsafe { Box::from_raw(req as *mut ModuleResponse) }; - let body = unsafe { CString::from_raw(boxed.body as *mut i8) }; - - // Explicitly calling drop here to feel good about myself - drop(body); - drop(boxed); +extern "C" fn cgi_cleanup(response: *const ModuleResponse) { + Response::cleanup(response); } |