From 348feff4cf0959d349287a7384ceb1d89a0543c4 Mon Sep 17 00:00:00 2001
From: bvisness <ben@handmade.network>
Date: Fri, 13 Jan 2023 18:52:31 +0000
Subject: [PATCH] Automatically handle HEAD requests (#104)

HEAD requests will be routed as GET requests, and Content-Type and Content-Length headers will always be sent for all requests.

Co-authored-by: Ben Visness <bvisness@gmail.com>
Reviewed-on: https://git.handmade.network/hmn/hmn/pulls/104
---
 src/website/requesthandling.go | 45 ++++++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/website/requesthandling.go b/src/website/requesthandling.go
index 70f2ed23..40390c01 100644
--- a/src/website/requesthandling.go
+++ b/src/website/requesthandling.go
@@ -12,6 +12,7 @@ import (
 	"net/url"
 	"path"
 	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -105,9 +106,14 @@ func (rb *RouteBuilder) Group(regex *regexp.Regexp, ms ...Middleware) RouteBuild
 }
 
 func (r *Router) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	method := req.Method
+	if method == http.MethodHead {
+		method = http.MethodGet // HEADs map to GETs for the purposes of routing
+	}
+
 nextroute:
 	for _, route := range r.Routes {
-		if route.Method != "" && req.Method != route.Method {
+		if route.Method != "" && method != route.Method {
 			continue
 		}
 
@@ -465,12 +471,36 @@ func doRequest(rw http.ResponseWriter, c *RequestContext, h Handler) {
 		}
 	}()
 
+	// Run the chosen handler
 	res := h(c)
 
 	if res.StatusCode == 0 {
 		res.StatusCode = http.StatusOK
 	}
 
+	// Set Content-Type and Content-Length if necessary. This behavior would in
+	// some cases be handled by http.ResponseWriter.Write, but we extract it so
+	// that HEAD requests always return both headers.
+
+	var preamble []byte // Any bytes we read to determine Content-Type
+	if res.Body != nil {
+		bodyLen := res.Body.Len()
+
+		if res.Header().Get("Content-Type") == "" {
+			preamble = res.Body.Next(512)
+			rw.Header().Set("Content-Type", http.DetectContentType(preamble))
+		}
+		if res.Header().Get("Content-Length") == "" {
+			rw.Header().Set("Content-Length", strconv.Itoa(bodyLen))
+		}
+	}
+
+	// Ensure we send no body for HEAD requests
+	if c.Req.Method == http.MethodHead {
+		res.Body = nil
+	}
+
+	// Send remaining response headers
 	for name, vals := range res.Header() {
 		for _, val := range vals {
 			rw.Header().Add(name, val)
@@ -478,7 +508,18 @@ func doRequest(rw http.ResponseWriter, c *RequestContext, h Handler) {
 	}
 	rw.WriteHeader(res.StatusCode)
 
+	// Send response body
 	if res.Body != nil {
-		io.Copy(rw, res.Body)
+		// Write preamble, if any
+		_, err := rw.Write(preamble)
+		if err != nil {
+			logging.Error().Err(err).Msg("Failed to write response preamble")
+		}
+
+		// Write remainder of body
+		_, err = io.Copy(rw, res.Body)
+		if err != nil {
+			logging.Error().Err(err).Msg("copied res.Body")
+		}
 	}
 }