add smalldog as the module ffi crate

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);
 }