improve saftey

1 files changed, 20 insertions, 45 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) {