handle HEAD correctly, fail with 405 when not GET/HEAD

1 files changed, 21 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index e352fe5..4809956 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -71,7 +71,7 @@ async fn handle_request(mut stream: TcpStream, caddr: SocketAddr) {
 		}
 
 		// Fall through the match. We should guarantee that, if we end up here,
-		// we got a Status::Complete from the parser
+		// we got a Status::Complete from the parser.
 		let response = form_response(req, caddr).await;
 
 		if let Err(err) = stream.write_all(&response).await {
@@ -105,26 +105,39 @@ fn get_header<'h>(
 		.transpose()
 }
 
+/// Create a response, whatever it is. This function forms responses for
+/// success, and errors.
 async fn form_response<'req>(request: Request<'req, 'req>, caddr: SocketAddr) -> Vec<u8> {
+	let is_head = match request.method.unwrap() {
+		"GET" => false,
+		"HEAD" => true,
+		// This is not strictly correct, but it's correct enough? :)
+		// 405 is "The request method is known by the server but
+		// is not supported by the target resource"
+		_ => return error(405, "we do not support that request method"),
+	};
+
 	let minor = request.version.unwrap();
 	let status = 200;
 	let message = "Gotcha";
 
-	let mut response = format!("HTTP/1.{minor} {status} {message}\n");
-
 	let requested_type = match get_header(request.headers, "content-type") {
 		Err(e) => return e,
 		Ok(None) | Ok(Some("text/plain")) => ContentType::Plain,
 		Ok(Some("application/json")) => ContentType::Json,
-		Ok(Some(_)) => return error(400, "unknown content-type"),
+		Ok(Some(_)) => return error(415, "unsupported content-type"),
 	};
 
+	// Make our response body!
 	let body = match requested_type {
 		ContentType::Plain => format!("{}", caddr.ip()),
 		ContentType::Json => format!(r#"{{"ip": "{}"}}"#, caddr.ip()),
 	};
 	let length = body.len();
 
+	// And now we form the request :D
+	let mut response = format!("HTTP/1.{minor} {status} {message}\n");
+
 	macro_rules! header {
 		($content:expr) => {
 			response.push_str(&format!($content));
@@ -136,8 +149,10 @@ async fn form_response<'req>(request: Request<'req, 'req>, caddr: SocketAddr) ->
 	header!("Content-Type: {requested_type}");
 	header!("Cache-Control: no-store");
 
-	response.push('\n');
-	response.push_str(&body);
+	if !is_head {
+		response.push('\n');
+		response.push_str(&body);
+	}
 
 	response.into_bytes()
 }