Add support for hiding request details

Refactor code
Use Viper to manage config both via flags and environment variables
More error handling of parsing templates, and return a response if an error.
Register middleware for prometheus metrics on error-page handler

go.mod

@@ -3,6 +3,9 @@ module github.com/181192/custom-error-pages
 go 1.15
 
 require (
+	github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c
 	github.com/prometheus/client_golang v1.7.1
 	github.com/rs/zerolog v1.19.0
+	github.com/spf13/pflag v1.0.5
+	github.com/spf13/viper v1.7.1
 )

main.go

@@ -1,30 +1,12 @@
 package main
 
 import (
-	"flag"
-	"fmt"
 	"net/http"
-	"os"
-	"strconv"
 
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"github.com/rs/zerolog"
+	"github.com/181192/custom-error-pages/pkg/handlers"
+	"github.com/181192/custom-error-pages/pkg/util"
 	"github.com/rs/zerolog/log"
-)
-
-const (
-	httpListenAddress    = "listen"
-	httpListenAddressEnv = "HTTP_LISTEN_ADDRESS"
-
-	debug    = "debug"
-	debugEnv = "DEBUG"
-
-	logColor    = "log-color"
-	logColorEnv = "LOG_COLOR"
-
-	errFilesPath    = "error-files-path"
-	errFilesPathEnv = "ERROR_FILES_PATH"
+	"github.com/spf13/viper"
 )
 
 var (
@@ -33,85 +15,18 @@ var (
 	date      = "unversioned"
 )
 
-// Options cli options
-type Options struct {
-	HTTPListenAddress string
-	Debug             bool
-	ColorLog          bool
-	ErrFilesPath      string
-}
-
 func main() {
+	util.InitFlags()
+	util.ConfigureLogger()
 
-	var opts Options
-
-	flag.BoolVar(&opts.Debug, debug, LookupEnvOrBool(debugEnv, false), "sets log level to debug")
-	flag.BoolVar(&opts.ColorLog, logColor, LookupEnvOrBool(logColorEnv, false), "sets log format to human-friendly, colorized output")
-	flag.StringVar(&opts.HTTPListenAddress, httpListenAddress, LookupEnvOrString(httpListenAddressEnv, ":8080"), "http server address")
-	flag.StringVar(&opts.ErrFilesPath, errFilesPath, LookupEnvOrString(errFilesPathEnv, "./www"), "the location on disk of files served by the handler")
-
-	flag.Parse()
-
-	if opts.ColorLog {
-		log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
-	}
-
-	zerolog.SetGlobalLevel(zerolog.InfoLevel)
-	if opts.Debug {
-		zerolog.SetGlobalLevel(zerolog.DebugLevel)
-	}
+	http.Handle("/", handlers.ErrorPage())
+	http.Handle("/metrics", handlers.Metrics())
+	http.Handle("/healthz", handlers.Health())
 
-	prometheus.Register(requestCount)
-	prometheus.Register(requestDuration)
-
-	http.Handle("/", &opts)
-	http.Handle("/metrics", promhttp.Handler())
-	http.HandleFunc("/healthz", healthHandler)
-
-	log.Debug().Msgf("config values: %+v", getConfig(flag.CommandLine))
+	httpListenAddress := viper.GetString(util.HTTPListenAddress)
+	log.Debug().Msgf("config values: %+v", viper.AllSettings())
 	log.Info().Msgf("version=%s, commit=%s, date=%s", version, gitCommit, date)
-	log.Info().Msgf("listening on %s", opts.HTTPListenAddress)
-	err := http.ListenAndServe(opts.HTTPListenAddress, nil)
+	log.Info().Msgf("listening on %s", httpListenAddress)
+	err := http.ListenAndServe(httpListenAddress, nil)
 	log.Fatal().Msg(err.Error())
 }
-
-// LookupEnvOrString lookup env from key or return default value
-func LookupEnvOrString(key string, defaultVal string) string {
-	if val, ok := os.LookupEnv(key); ok {
-		return val
-	}
-	return defaultVal
-}
-
-// LookupEnvOrInt lookup env from key or return default value
-func LookupEnvOrInt(key string, defaultVal int) int {
-	if val, ok := os.LookupEnv(key); ok {
-		v, err := strconv.Atoi(val)
-		if err != nil {
-			log.Fatal().Msgf("LookupEnvOrInt[%s]: %v", key, err)
-		}
-		return v
-	}
-	return defaultVal
-}
-
-// LookupEnvOrBool lookup env from key or return default value
-func LookupEnvOrBool(key string, defaultVal bool) bool {
-	if val, ok := os.LookupEnv(key); ok {
-		v, err := strconv.ParseBool(val)
-		if err != nil {
-			log.Fatal().Msgf("LookupEnvOrBool[%s]: %v", key, err)
-		}
-		return v
-	}
-	return defaultVal
-}
-
-func getConfig(fs *flag.FlagSet) []string {
-	cfg := make([]string, 0, 10)
-	fs.VisitAll(func(f *flag.Flag) {
-		cfg = append(cfg, fmt.Sprintf("%s:%q", f.Name, f.Value.String()))
-	})
-
-	return cfg
-}

error.go → pkg/handlers/error-page.go

@@ -1,16 +1,17 @@
-package main
+package handlers
 
 import (
 	"encoding/json"
-	"fmt"
 	"net/http"
-	"os"
 	"strconv"
 	"strings"
 	"text/template"
-	"time"
 
+	"github.com/181192/custom-error-pages/pkg/metrics"
+	"github.com/181192/custom-error-pages/pkg/util"
+	"github.com/oxtoacart/bpool"
 	"github.com/rs/zerolog/log"
+	"github.com/spf13/viper"
 )
 
 const (
@@ -48,11 +49,17 @@ const (
 	HTML = "text/html"
 )
 
+var bufpool *bpool.BufferPool
+
+func init() {
+	bufpool = bpool.NewBufferPool(64)
+}
+
 type errorPageData struct {
-	Code    string               `json:"code"`
-	Title   string               `json:"title"`
-	Message []string             `json:"message"`
-	Details errorPageDataDetails `json:"details,omitempty"`
+	Code     string                `json:"code"`
+	Title    string                `json:"title"`
+	Messages []string              `json:"messages"`
+	Details  *errorPageDataDetails `json:"details,omitempty"`
 }
 
 type errorPageDataDetails struct {
@@ -64,24 +71,28 @@ type errorPageDataDetails struct {
 	RequestID   string `json:"requestId"`
 }
 
-func newErrorPageData(req *http.Request, message []string) errorPageData {
-	statusCode := req.Header.Get(CodeHeader)
-	statusCodeNumber, _ := strconv.Atoi(req.Header.Get(CodeHeader))
-	statusText := http.StatusText(statusCodeNumber)
+func newErrorPageData(req *http.Request, code int, messages []string) errorPageData {
+	title := http.StatusText(code)
 
-	return errorPageData{
-		Code:    statusCode,
-		Title:   statusText,
-		Message: message,
-		Details: errorPageDataDetails{
+	data := errorPageData{
+		Code:     strconv.Itoa(code),
+		Title:    title,
+		Messages: messages,
+	}
+
+	hideDetails := viper.GetBool(util.HideDetails)
+	if !hideDetails {
+		data.Details = &errorPageDataDetails{
 			OriginalURI: req.Header.Get(OriginalURI),
 			Namespace:   req.Header.Get(Namespace),
 			IngressName: req.Header.Get(IngressName),
 			ServiceName: req.Header.Get(ServiceName),
 			ServicePort: req.Header.Get(ServicePort),
 			RequestID:   req.Header.Get(RequestID),
-		},
+		}
 	}
+
+	return data
 }
 
 func getFormat(req *http.Request) string {
@@ -114,7 +125,7 @@ func getStatusCode(req *http.Request) int {
 	return code
 }
 
-func getMessage(code int) []string {
+func getMessages(code int) []string {
 	switch code {
 	case http.StatusNotFound:
 		return []string{"The page you're looking for could not be found."}
@@ -125,61 +136,79 @@ func getMessage(code int) []string {
 	}
 }
 
-// HTMLResponse returns html reponse
-func HTMLResponse(w http.ResponseWriter, r *http.Request, path string) {
+// htmlResponse returns html reponse
+func htmlResponse(w http.ResponseWriter, r *http.Request) {
 	code := getStatusCode(r)
-	message := getMessage(code)
+	messages := getMessages(code)
 
-	stylesPath := fmt.Sprintf("%v/styles.css", path)
-	styles, err := os.Open(stylesPath)
+	buf := bufpool.Get()
+	defer bufpool.Put(buf)
 
-	file := fmt.Sprintf("%v/template.html", path)
-	f, err := os.Open(file)
+	templatesDir := viper.GetString(util.ErrFilesPath) + "/*"
+	templates, err := template.ParseGlob(templatesDir)
 	if err != nil {
-		log.Warn().Msgf("unexpected error opening file: %v", err)
-		JSONResponse(w, r)
+		log.Error().Msgf("Failed to parse template %s", err)
+		w.Header().Set(ContentType, JSON)
+		w.WriteHeader(http.StatusInternalServerError)
+		body, _ := json.Marshal(newErrorPageData(r, http.StatusInternalServerError,
+			[]string{"Ups, this should not have happened", "Failed to parse templates"}))
+		w.Write(body)
 		return
 	}
-	defer f.Close()
 
-	log.Debug().Msgf("serving custom error response for code %v and format %v from file %v", code, HTML, file)
-	tmpl := template.Must(template.ParseFiles(f.Name(), styles.Name()))
+	data := newErrorPageData(r, code, messages)
+	err = templates.ExecuteTemplate(buf, "index", data)
+	if err != nil {
+		log.Error().Msgf("Failed to execute template %s", err)
+		w.Header().Set(ContentType, JSON)
+		w.WriteHeader(http.StatusInternalServerError)
+		body, _ := json.Marshal(newErrorPageData(r, http.StatusInternalServerError,
+			[]string{"Ups, this should not have happened", "Failed to execute templates"}))
+		w.Write(body)
+		return
+	}
 
 	w.Header().Set(ContentType, HTML)
 	w.WriteHeader(code)
-
-	data := newErrorPageData(r, message)
-	tmpl.Execute(w, data)
+	buf.WriteTo(w)
 }
 
-// JSONResponse returns json reponse
-func JSONResponse(w http.ResponseWriter, r *http.Request) {
+// jsonResponse returns json reponse
+func jsonResponse(w http.ResponseWriter, r *http.Request) {
 	code := getStatusCode(r)
-	message := getMessage(code)
+	message := getMessages(code)
+
+	body, err := json.Marshal(newErrorPageData(r, code, message))
+	if err != nil {
+		log.Error().Msgf("Failed to marshal json response %s", err)
+		w.Header().Set(ContentType, JSON)
+		w.WriteHeader(http.StatusInternalServerError)
+		body, _ := json.Marshal(newErrorPageData(r, http.StatusInternalServerError,
+			[]string{"Ups, this should not have happened", "Failed to marshal json response"}))
+		w.Write(body)
+		return
+	}
 
 	w.Header().Set(ContentType, JSON)
 	w.WriteHeader(code)
-	body, _ := json.Marshal(newErrorPageData(r, message))
 	w.Write(body)
 }
 
-// ServeHttp error handler
-func (opts *Options) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	start := time.Now()
-	format := getFormat(r)
-
-	switch format {
-	case JSON:
-		JSONResponse(w, r)
-	default:
-		HTMLResponse(w, r, opts.ErrFilesPath)
-	}
-
-	duration := time.Now().Sub(start).Seconds()
+// ErrorPage error handler
+func ErrorPage() http.Handler {
+	return metrics.Measure(errorPage())
+}
 
-	proto := strconv.Itoa(r.ProtoMajor)
-	proto = fmt.Sprintf("%s.%s", proto, strconv.Itoa(r.ProtoMinor))
+func errorPage() http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		format := getFormat(r)
 
-	requestCount.WithLabelValues(proto).Inc()
-	requestDuration.WithLabelValues(proto).Observe(duration)
+		switch format {
+		case JSON:
+			jsonResponse(w, r)
+		default:
+			htmlResponse(w, r)
+		}
+	}
+	return http.HandlerFunc(fn)
 }