begin big refactor

Author: nullbio

abcmiddleware/errors.go

@@ -128,7 +128,7 @@ func (m *errorManager) Errors(ctrl AppHandler) http.HandlerFunc {
 		}
 
 		// Get the Request ID scoped logger
-		log := Log(r)
+		log := Logger(r)
 
 		fields := []zapcore.Field{
 			zap.String("method", r.Method),

abcmiddleware/log.go

@@ -1,103 +1,73 @@
 package abcmiddleware
 
 import (
-	"context"
 	"fmt"
 	"net/http"
 	"time"
 
-	chimiddleware "github.com/go-chi/chi/middleware"
 	"go.uber.org/zap"
 )
 
-// zapResponseWriter is a wrapper that includes that http status and size for logging
-type zapResponseWriter struct {
-	http.ResponseWriter
-	status int
-	size   int
+type zapLogMiddleware struct {
+	logger *zap.Logger
 }
 
-// Zap middleware handles web request logging using Zap
-func (m Middleware) Zap(next http.Handler) http.Handler {
-	fn := func(w http.ResponseWriter, r *http.Request) {
-		t := time.Now()
-		zw := &zapResponseWriter{ResponseWriter: w}
-
-		// Serve the request
-		next.ServeHTTP(zw, r)
+// ZapLog returns a logging middleware that outputs details about a request
+func ZapLog(logger *zap.Logger) MW {
+	return zapLogMiddleware{logger: logger}
+}
 
-		// Write the request log line
-		writeZap(m.Log, r, t, zw.status, zw.size)
-	}
+// Zap middleware handles web request logging using Zap
+func (z zapLogMiddleware) Wrap(next http.Handler) http.Handler {
+	return zapLogger{mid: z, next: next}
+}
 
-	return http.HandlerFunc(fn)
+type zapLogger struct {
+	mid  zapLogMiddleware
+	next http.Handler
 }
 
-// RequestIDLogger middleware creates a derived logger to include logging of the
-// Request ID, and inserts it into the context object
-func (m Middleware) RequestIDLogger(next http.Handler) http.Handler {
-	fn := func(w http.ResponseWriter, r *http.Request) {
-		requestID := chimiddleware.GetReqID(r.Context())
-		derivedLogger := m.Log.With(zap.String("request_id", requestID))
-		r = r.WithContext(context.WithValue(r.Context(), CtxLoggerKey, derivedLogger))
-		next.ServeHTTP(w, r)
-	}
+func (z zapLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	startTime := time.Now()
+	zw := &zapResponseWriter{ResponseWriter: w}
 
-	return http.HandlerFunc(fn)
-}
+	// Serve the request
+	z.next.ServeHTTP(zw, r)
 
-// Log returns the Request ID scoped logger from the request Context
-// and panics if it cannot be found. This function is only ever used
-// by your controllers if your app uses the RequestID middlewares,
-// otherwise you should use the controller's receiver logger directly.
-func Log(r *http.Request) *zap.Logger {
-	v := r.Context().Value(CtxLoggerKey)
-	log, ok := v.(*zap.Logger)
-	if !ok {
-		panic("cannot get derived request id logger from context object")
-	}
-	return log
+	// Write the request log line
+	z.writeZap(zw, r, startTime)
 }
 
-func writeZap(log *zap.Logger, r *http.Request, t time.Time, status int, size int) {
-	elapsed := time.Now().Sub(t)
+func (z zapLogger) writeZap(zw *zapResponseWriter, r *http.Request, startTime time.Time) {
+	elapsed := time.Now().Sub(startTime)
 	var protocol string
 	if r.TLS == nil {
 		protocol = "http"
 	} else {
 		protocol = "https"
 	}
 
-	v := r.Context().Value(CtxLoggerKey)
+	logger := z.mid.logger
+	v := r.Context().Value(CTXKeyLogger)
 	if v != nil {
 		var ok bool
-		log, ok = v.(*zap.Logger)
+		logger, ok = v.(*zap.Logger)
 		if !ok {
 			panic("cannot get derived request id logger from context object")
 		}
 	}
 
 	// log all the fields
-	log.Info(fmt.Sprintf("%s request", protocol),
-		zap.Int("status", status),
+	logger.Info(fmt.Sprintf("%s request", protocol),
+		zap.Int("status", zw.status),
+		zap.Int("size", zw.size),
+		zap.Bool("hijacked", zw.hijacked),
 		zap.String("method", r.Method),
 		zap.String("uri", r.RequestURI),
 		zap.Bool("tls", r.TLS != nil),
 		zap.String("protocol", r.Proto),
 		zap.String("host", r.Host),
 		zap.String("remote_addr", r.RemoteAddr),
-		zap.Int("size", size),
 		zap.Duration("elapsed", elapsed),
 	)
 }
-
-func (z *zapResponseWriter) WriteHeader(code int) {
-	z.status = code
-	z.ResponseWriter.WriteHeader(code)
-}
-
-func (z *zapResponseWriter) Write(b []byte) (int, error) {
-	size, err := z.ResponseWriter.Write(b)
-	z.size += size
-	return size, err
-}

abcmiddleware/middleware.go

@@ -6,9 +6,6 @@ import (
 	"go.uber.org/zap"
 )
 
-// CtxLoggerKey is the http.Request Context lookup key for the request ID logger
-const CtxLoggerKey = "request_id_logger"
-
 // MiddlewareFunc is the function signature for Chi's Use() middleware
 type MiddlewareFunc func(http.Handler) http.Handler
 
@@ -18,3 +15,8 @@ type Middleware struct {
 	// create a derived logger that includes the request ID.
 	Log *zap.Logger
 }
+
+// MW is an interface defining middleware wrapping
+type MW interface {
+	Wrap(http.Handler) http.Handler
+}

abcmiddleware/recover.go

@@ -7,57 +7,79 @@ import (
 	"go.uber.org/zap"
 )
 
-// Recover middleware recovers panics that occur and gracefully logs their error
-func (m Middleware) Recover(next http.Handler) http.Handler {
-	fn := func(w http.ResponseWriter, r *http.Request) {
-		defer func() {
-			if err := recover(); err != nil {
-				var protocol string
-				if r.TLS == nil {
-					protocol = "http"
-				} else {
-					protocol = "https"
-				}
-
-				var log *zap.Logger
-				v := r.Context().Value(CtxLoggerKey)
-				if v != nil {
-					var ok bool
-					log, ok = v.(*zap.Logger)
-					if !ok {
-						panic("cannot get derived request id logger from context object")
-					}
-					// log with the request_id scoped logger
-					log.Error(fmt.Sprintf("%s request error", protocol),
-						zap.String("method", r.Method),
-						zap.String("uri", r.RequestURI),
-						zap.Bool("tls", r.TLS != nil),
-						zap.String("protocol", r.Proto),
-						zap.String("host", r.Host),
-						zap.String("remote_addr", r.RemoteAddr),
-						zap.String("error", fmt.Sprintf("%+v", err)),
-					)
-				} else {
-					// log with the logger attached to middleware struct if
-					// cannot find request_id scoped logger
-					m.Log.Error(fmt.Sprintf("%s request error", protocol),
-						zap.String("method", r.Method),
-						zap.String("uri", r.RequestURI),
-						zap.Bool("tls", r.TLS != nil),
-						zap.String("protocol", r.Proto),
-						zap.String("host", r.Host),
-						zap.String("remote_addr", r.RemoteAddr),
-						zap.String("error", fmt.Sprintf("%+v", err)),
-					)
-				}
-
-				// Return a http 500 with the HTTP body of "Internal Server Error"
-				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
-			}
-		}()
-
-		next.ServeHTTP(w, r)
+// ZapRecover will attempt to log a panic, as well as produce a reasonable
+// error for the client by calling the passed in errorHandler function.
+//
+// It uses the zap logger and attempts to look up a request-scoped logger
+// created with this package before using the passed in logger.
+//
+// The zap logger that's used here should be careful to enable stacktrace
+// logging for any levels that they require it for.
+func ZapRecover(fallback *zap.Logger, errorHandler http.HandlerFunc) MW {
+	return zapRecoverMiddleware{
+		fallback: fallback,
+		eh:       errorHandler,
+	}
+}
+
+type zapRecoverMiddleware struct {
+	fallback *zap.Logger
+	eh       http.HandlerFunc
+}
+
+func (z zapRecoverMiddleware) Wrap(next http.Handler) http.Handler {
+	return zapRecoverer{
+		zr:   z,
+		next: next,
+	}
+}
+
+type zapRecoverer struct {
+	zr   zapRecoverMiddleware
+	next http.Handler
+}
+
+// recoverPanic was mostly adapted from abcweb
+func (z zapRecoverer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	defer z.recoverNicely(w, r)
+	z.next.ServeHTTP(w, r)
+}
+
+func (z zapRecoverer) recoverNicely(w http.ResponseWriter, r *http.Request) {
+	err := recover()
+	if err == nil {
+		return
+	}
+
+	var protocol string
+	if r.TLS == nil {
+		protocol = "http"
+	} else {
+		protocol = "https"
+	}
+
+	if z.zr.eh != nil {
+		z.zr.eh(w, r)
+	}
+
+	logger := z.zr.fallback
+	v := r.Context().Value(CTXKeyLogger)
+	if v != nil {
+		var ok bool
+		logger, ok = v.(*zap.Logger)
+		if !ok {
+			panic("cannot get derived request id logger from context object")
+		}
 	}
 
-	return http.HandlerFunc(fn)
+	logger.Error(fmt.Sprintf("%s request error", protocol),
+		zap.String("method", r.Method),
+		zap.String("uri", r.RequestURI),
+		zap.Bool("tls", r.TLS != nil),
+		zap.String("protocol", r.Proto),
+		zap.String("host", r.Host),
+		zap.String("remote_addr", r.RemoteAddr),
+		zap.String("panic", fmt.Sprintf("%+v", err)),
+		zap.Stack("stacktrace"),
+	)
 }

abcmiddleware/request_id.go

@@ -0,0 +1,79 @@
+package abcmiddleware
+
+import (
+	"context"
+	"net/http"
+
+	chimiddleware "github.com/go-chi/chi/middleware"
+	"go.uber.org/zap"
+)
+
+type ctxKey int
+
+const (
+	// CTXKeyLogger is the key under which the request scoped logger is placed
+	CTXKeyLogger ctxKey = iota
+)
+
+// RequestIDHeader sets the X-Request-ID header to the chi request id
+// This must be used after the chi request id middleware.
+func RequestIDHeader(next http.Handler) http.Handler {
+	return reqIDInserter{next: next}
+}
+
+type reqIDInserter struct {
+	next http.Handler
+}
+
+func (re reqIDInserter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if reqID := chimiddleware.GetReqID(r.Context()); len(reqID) != 0 {
+		w.Header().Set("X-Request-ID", reqID)
+	}
+	re.next.ServeHTTP(w, r)
+}
+
+// ZapRequestIDLogger returns a request id logger middleware. This only works
+// if chi has inserted a request id into the stack first.
+func ZapRequestIDLogger(logger *zap.Logger) MW {
+	return zapReqLoggerMiddleware{logger: logger}
+}
+
+type zapReqLoggerMiddleware struct {
+	logger *zap.Logger
+}
+
+func (z zapReqLoggerMiddleware) Wrap(next http.Handler) http.Handler {
+	return zapReqLoggerInserter{logger: z.logger, next: next}
+}
+
+type zapReqLoggerInserter struct {
+	logger *zap.Logger
+	next   http.Handler
+}
+
+func (z zapReqLoggerInserter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	requestID := chimiddleware.GetReqID(r.Context())
+
+	derivedLogger := z.logger.With(zap.String("request_id", requestID))
+
+	r = r.WithContext(context.WithValue(r.Context(), CTXKeyLogger, derivedLogger))
+	z.next.ServeHTTP(w, r)
+}
+
+// Logger returns the Request ID scoped logger from the request Context
+// and panics if it cannot be found. This function is only ever used
+// by your controllers if your app uses the RequestID middlewares,
+// otherwise you should use the controller's receiver logger directly.
+func Logger(r *http.Request) *zap.Logger {
+	return LoggerCTX(r.Context())
+}
+
+// LoggerCTX retrieves a logger from a context.
+func LoggerCTX(ctx context.Context) *zap.Logger {
+	v := ctx.Value(CTXKeyLogger)
+	log, ok := v.(*zap.Logger)
+	if !ok {
+		panic("cannot get derived request id logger from context object")
+	}
+	return log
+}