light refactor

2 files changed, 83 insertions, 42 deletions

diff --git a/README.md b/README.md
index 3b6b0df..9b51bfa 100644
--- a/README.md
+++ b/README.md
@@ -16,4 +16,10 @@ Sets the following environmental variables for the CGI script, many following [R
 - **`QUERY_STRING`** to the query part of the URI
 - **`REQUEST_METHOD`** to the HTTP request method
 
+Additionally, corgi will set environment variables for the HTTP request headers.
+They will be uppercased and hyphens replaced with underscores.
+
+Any environmental variable may be overridden if it is set in the
+configuration file.
+
 [rfc]: https://datatracker.ietf.org/doc/html/rfc3875
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index cddcab6..c83c0bd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,7 @@ use std::{net::SocketAddr, pin::Pin};
 use confindent::Confindent;
 use http_body_util::{BodyExt, Full};
 use hyper::{
-	Request, Response, StatusCode,
+	HeaderMap, Request, Response, StatusCode,
 	body::{Body, Bytes, Incoming},
 	server::conn::http1,
 	service::Service,
@@ -92,22 +92,11 @@ impl Svc {
 			.env("QUERY_STRING", query)
 			.env("REQUEST_METHOD", method);
 
-		for (header, value) in headers {
-			if let Some(header) = header {
-				let hname = header.as_str();
-				let hvalue = value.to_str().unwrap();
-				cmd.env(hname.to_ascii_uppercase(), hvalue);
-
-				if hname.to_ascii_lowercase() == "user-agent" {
-					println!("USER_AGENT: {hvalue}");
-				}
-
-				if hname.to_ascii_lowercase() == "host" {
-					println!("HOST: {hvalue}");
-				}
-			}
-		}
+		// Set env associated with the HTTP request headers
+		Self::set_http_env(headers, &mut cmd);
 
+		// Set env specified in the conf. Be sure we do this after we
+		// set the HTTP headers as to overwrite any we might want
 		for (key, value) in &settings.env {
 			cmd.env(key.to_ascii_uppercase(), value);
 		}
@@ -116,56 +105,102 @@ impl Svc {
 			cmd.env("CONTENT_LENGTH", content_length.to_string());
 		}*/
 
+		let cgi_response = Self::call_and_parse_cgi(cmd).await;
+
+		let mut response =
+			Response::builder().status(StatusCode::from_u16(cgi_response.status).unwrap());
+
+		for (key, value) in cgi_response.headers {
+			response = response.header(key, value);
+		}
+
+		response.body(Full::new(Bytes::from(body.to_vec()))).unwrap()
+	}
+
+	fn make<B: Into<Bytes>>(b: B) -> Response<Full<Bytes>> {
+		Response::builder().body(Full::new(b.into())).unwrap()
+	}
+
+	fn set_http_env(headers: HeaderMap, cmd: &mut Command) {
+		for (key, value) in headers.iter() {
+			let key_str = key.as_str();
+
+			let mut key_upper = String::with_capacity(key_str.len() + 5);
+			key_upper.push_str("HTTP_");
+
+			for ch in key_str.chars() {
+				match ch {
+					_ if ch as u8 > 0x60 && ch as u8 <= 0x7A => {
+						key_upper.push((ch as u8 - 0x20) as char);
+					}
+					'-' => key_upper.push('_'),
+					ch => key_upper.push(ch),
+				}
+			}
+
+			match value.to_str() {
+				Ok(val_str) => {
+					cmd.env(key_upper, val_str);
+				}
+				Err(err) => {
+					eprintln!("value for header {key_str} is not a string: {err}")
+				}
+			}
+		}
+	}
+
+	async fn call_and_parse_cgi(mut cmd: Command) -> CgiResponse {
+		let mut response = CgiResponse {
+			// Default status code is 200 per RFC
+			status: 200,
+			headers: vec![],
+			body: None,
+		};
+
 		let output = cmd.output().await.unwrap();
 		let response_raw = output.stdout;
-		let mut response = Response::builder();
-
-		println!("{}", String::from_utf8_lossy(&response_raw));
 
 		let mut curr = response_raw.as_slice();
-		let mut status = None;
-		let mut headers = vec![];
-		let body = loop {
+		loop {
+			// Find the newline to know where this header ends
 			let nl = curr.iter().position(|b| *b == b'\n').expect("no nl in header");
 			let line = &curr[..nl];
 
+			// Find the colon to separate the key from the value
 			let colon = line.iter().position(|b| *b == b':').expect("no colon in header");
 			let key = &line[..colon];
 			let value = &line[colon + 1..];
-			headers.push((key, value));
+			response.headers.push((key.to_vec(), value.to_vec()));
 
 			let key_string = String::from_utf8_lossy(key);
 			if key_string == "Status" {
 				let value_string = String::from_utf8_lossy(value);
-				if let Some((raw_code, raw_msg)) = value_string.trim().split_once(' ') {
-					let code: usize = raw_code.parse().unwrap();
-					status = Some((code, raw_msg.to_owned()));
+				if let Some((raw_code, _raw_msg)) = value_string.trim().split_once(' ') {
+					let code: u16 = raw_code.parse().unwrap();
+					response.status = code;
 				}
 			}
 
 			// Body next
 			if curr[nl + 1] == b'\n' || (curr[nl + 1] == b'\r' && curr[nl + 2] == b'\n') {
-				break &curr[nl + 2..];
+				let body = &curr[nl + 2..];
+				if body.len() > 0 {
+					response.body = Some(body.to_vec());
+				}
 			} else {
 				curr = &curr[nl + 1..];
 			}
-		};
-
-		match status {
-			None => response = response.status(StatusCode::OK),
-			Some((code, _status)) => {
-				response = response.status(StatusCode::from_u16(code as u16).unwrap());
-			}
-		}
-
-		for (key, value) in headers {
-			response = response.header(key.to_vec(), value.to_vec());
 		}
 
-		response.body(Full::new(Bytes::from(body.to_vec()))).unwrap()
+		response
 	}
+}
 
-	fn make<B: Into<Bytes>>(b: B) -> Response<Full<Bytes>> {
-		Response::builder().body(Full::new(b.into())).unwrap()
-	}
+struct CgiResponse {
+	/// The Status header of the CGI response
+	status: u16,
+	/// Headers except "Status"
+	headers: Vec<(Vec<u8>, Vec<u8>)>,
+	/// CGI response body
+	body: Option<Vec<u8>>,
 }