grim/goserve

Gross fix, but this will set the content type when we're serving up the custom 404 page
draft default tip
2020-01-04, Gary Kramlich
0750e24e8e5c
Gross fix, but this will set the content type when we're serving up the custom 404 page
package main
import (
"fmt"
"io/ioutil"
"net/http"
)
type responseWriter struct {
w http.ResponseWriter
wroteHeader bool
code int
wrote int64
notFoundFile string
}
var _ http.ResponseWriter = (*responseWriter)(nil)
func (w *responseWriter) Header() http.Header {
return w.w.Header()
}
func (w *responseWriter) Write(data []byte) (int, error) {
if !w.wroteHeader {
w.wroteHeader = true
w.code = http.StatusOK
}
if w.code == http.StatusNotFound && w.notFoundFile != "" {
notFoundData, err := ioutil.ReadFile(w.notFoundFile)
if err != nil {
fmt.Printf(
"failed to load custom 404 file %q\n",
err,
)
} else {
// if the file loaded overwrite the output that the original
// handler gave us.
data = notFoundData
}
}
l, err := w.w.Write(data)
w.wrote += int64(l)
return l, err
}
func (w *responseWriter) WriteHeader(code int) {
if !w.wroteHeader {
w.wroteHeader = true
if code == http.StatusNotFound && w.notFoundFile != "" {
w.w.Header().Set("Content-Type", "text/html")
}
w.code = code
}
w.w.WriteHeader(w.code)
}
func (w *responseWriter) StatusCode() int {
return w.code
}
func (w *responseWriter) Written() int64 {
return w.wrote
}