2 files changed, 28 insertions, 57 deletions

diff --git a/smalldog/src/lib.rs b/smalldog/src/lib.rs
index e1ada11..adfe9c1 100644
--- a/smalldog/src/lib.rs
+++ b/smalldog/src/lib.rs
@@ -21,7 +21,8 @@ pub struct Request<'req> {
 }
 
 impl<'req> Request<'req> {
-	pub fn from_mod_request(request: *const ModuleRequest) -> Self {
+	pub fn from_mod_request(request: *const ModuleRequest<'req>) -> Self {
+		// SAFTEY: corgi will never give us a null pointer
 		let reqref = unsafe { request.as_ref() }.unwrap();
 
 		let mut headers = vec![];
@@ -61,86 +62,60 @@ pub struct ModuleResponse {
 	pub body: *const u8,
 }
 
-pub static mut HEADERS: [[*const ffi::c_char; 2]; 64] = [[ptr::null(), ptr::null()]; 64];
+const HEADERS_LEN: usize = 64;
+static mut HEADERS: [[*const ffi::c_char; 2]; HEADERS_LEN] = [[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,
+	headers: Vec<(Cow<'static, CStr>, Cow<'static, CStr>)>,
+	body: Vec<u8>,
 }
 
 impl Response {
 	pub fn new() -> Self {
 		Self {
 			headers: vec![],
-			body: ResponseBody::Building(String::new()),
+			body: vec![],
 		}
 	}
 
 	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() {
+		for (idx, (key, value)) in this.headers.iter().enumerate().take(HEADERS_LEN) {
 			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 headers_len = this.headers.len().min(HEADERS_LEN) as u64;
 		let boxed = Box::new(ModuleResponse {
 			status,
 			headers_len,
 			headers: unsafe { &HEADERS[..headers_len as usize] },
-			body_len,
-			body,
+			body_len: this.body.len() as u64,
+			body: this.body.as_ptr(),
 		});
 
 		Box::<ModuleResponse>::into_raw(boxed)
 	}
 
-	pub fn header<K: Into<CString>, V: Into<CString>>(&mut self, key: K, value: V) -> &mut Self {
+	pub fn header<K: Into<Cow<'static, CStr>>, V: Into<Cow<'static, CStr>>>(
+		&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 body(&mut self, vec: Vec<u8>) -> &mut Self {
+		self.body = vec;
+		self
 	}
 
 	pub fn cleanup(response: *const ModuleResponse) {
diff --git a/stats_module/src/lib.rs b/stats_module/src/lib.rs
index 96cbaa9..a56d22d 100644
--- a/stats_module/src/lib.rs
+++ b/stats_module/src/lib.rs
@@ -1,5 +1,3 @@
-use std::ffi::CStr;
-
 use rusqlite::{Connection, params};
 use smalldog::{ModuleRequest, ModuleResponse, Request, Response};
 use time::{Duration, OffsetDateTime};
@@ -7,6 +5,7 @@ use time::{Duration, OffsetDateTime};
 #[unsafe(no_mangle)]
 extern "C" fn cgi_handle(req: *const ModuleRequest) -> *const ModuleResponse {
 	let mut response = Response::new();
+	let mut body = String::new();
 
 	let request = Request::from_mod_request(req);
 	let db = if let Some(path) = request.header("CORGI_STATS_DB") {
@@ -32,27 +31,24 @@ extern "C" fn cgi_handle(req: *const ModuleRequest) -> *const ModuleResponse {
 
 	agents.sort_by(|a, b| a.0.cmp(&b.0).reverse());
 
-	response.push_str("<p>In the last fifteen minutes:<br/><code><pre>");
-	response.push_str("total | req/m | agent\n");
+	body.push_str("<p>In the last fifteen minutes:<br/><code><pre>");
+	body.push_str("total | req/m | agent\n");
 	for (count, agent) in &agents {
-		response.push_str(&format!(
+		body.push_str(&format!(
 			"{count:<5} | {:<5.1} | {agent}\n",
 			*count as f32 / 15.0
 		));
 	}
-	response.push_str("</pre></code></p>");
+	body.push_str("</pre></code></p>");
+	response.body(body.into_bytes());
 
 	response.into_mod_response(200)
 }
 
 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 mut response = Response::new();
-	response.push_str(msg.as_ref());
+	response.header(c"Content-Length", c"text/html");
+	response.body(msg.as_ref().as_bytes().to_vec());
 
 	response.into_mod_response(code)
 }