Add support for custom status codes in "serveFile". (#128)

* Add support for status codes in "serveFile".

* Prevent having to create an extra buffer and misusing memory.

* Avoid overwriting status code if one wasn't passed. Remove unnecessary log statement.

internal/server/handlers.go

@@ -63,7 +63,7 @@ func (s *Server) showOrRender(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// If the path is not a directory, then it's a file, so we can render it
-	s.serveFile(currentPath, w, r)
+	s.serveFile(0, currentPath, w, r)
 }
 
 func (s *Server) walk(requestedPath string, w http.ResponseWriter, r *http.Request) {
@@ -72,7 +72,7 @@ func (s *Server) walk(requestedPath string, w http.ResponseWriter, r *http.Reque
 	for _, index := range []string{"index.html", "index.htm"} {
 		indexPath := filepath.Join(requestedPath, index)
 		if _, err := os.Stat(indexPath); err == nil {
-			s.serveFile(indexPath, w, r)
+			s.serveFile(0, indexPath, w, r)
 			return
 		}
 	}
@@ -165,11 +165,30 @@ func (s *Server) walk(requestedPath string, w http.ResponseWriter, r *http.Reque
 	}
 }
 
+// statusCodeHijacker is a response writer that captures the status code
+// and the body of the response that would have been sent to the client.
+type statusCodeHijacker struct {
+	http.ResponseWriter
+	givenStatusCode int
+}
+
+// WriteHeader captures the status code that would have been sent to the client.
+func (s *statusCodeHijacker) WriteHeader(code int) {
+	s.givenStatusCode = code
+}
+
 // serveFile serves a file with the appropriate headers, including support
 // for ETag and Last-Modified headers, as well as range requests.
-func (s *Server) serveFile(fp string, w http.ResponseWriter, r *http.Request) {
-	f, err := os.Open(fp)
+// If the status code is not 0, the status code provided will be used
+// when serving the file in the given path.
+func (s *Server) serveFile(statusCode int, location string, w http.ResponseWriter, r *http.Request) {
+	f, err := os.Open(location)
 	if err != nil {
+		if os.IsNotExist(err) {
+			httpError(http.StatusNotFound, w, "404 not found")
+			return
+		}
+
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
@@ -182,7 +201,7 @@ func (s *Server) serveFile(fp string, w http.ResponseWriter, r *http.Request) {
 	}
 
 	var ctype string
-	if local := getContentTypeForFilename(filepath.Base(fp)); local != "" {
+	if local := getContentTypeForFilename(filepath.Base(location)); local != "" {
 		ctype = local
 	}
 
@@ -218,7 +237,19 @@ func (s *Server) serveFile(fp string, w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", ctype)
 	}
 
-	http.ServeContent(w, r, fi.Name(), fi.ModTime(), f)
+	// Check if the caller changed the status code, if not, simply call
+	// the appropriate handler/
+	if statusCode == 0 {
+		http.ServeContent(w, r, fi.Name(), fi.ModTime(), f)
+		return
+	}
+
+	// Write the status code sent by the caller.
+	w.WriteHeader(statusCode)
+
+	// Call serve content with the hijacked response writer, which won't
+	// be able to overwrite the status code.
+	http.ServeContent(&statusCodeHijacker{ResponseWriter: w}, r, fi.Name(), fi.ModTime(), f)
 }
 
 // healthCheck is a simple health check endpoint that returns 200 OK

internal/server/startup.go

@@ -65,9 +65,7 @@ func (s *Server) PrintStartup() {
 	}
 
 	if !s.DisableRedirects {
-		if s.redirects == nil {
-			fmt.Fprintf(s.LogOutput, "%s Redirections enabled but no redirections found in %q\n", startupPrefix, s.getPathToRedirectionsFile())
-		} else {
+		if s.redirects != nil {
 			fmt.Fprintf(s.LogOutput, "%s Redirections enabled from %q (found %d redirections)\n", startupPrefix, s.getPathToRedirectionsFile(), len(s.redirects.Rules))
 		}
 	}