Improve response logger and fix '/' route

polyculeconnect/middlewares/logger.go

@@ -5,25 +5,46 @@ import (
 	"net/http"
 	"time"
 
-	"git.faercol.me/faercol/polyculeconnect/polyculeconnect/helpers"
 	"github.com/sirupsen/logrus"
 )
 
-var defaultResponseInfo = helpers.ResponseInfo{
-	ReturnCode:    -1,
-	ContentLength: -1,
+type loggedResponseWriter struct {
+	w             http.ResponseWriter
+	statusCode    int
+	contentLength int
+}
+
+func (lr *loggedResponseWriter) Header() http.Header {
+	return lr.w.Header()
+}
+
+func (lr *loggedResponseWriter) Write(dat []byte) (int, error) {
+	if lr.statusCode < 100 {
+		lr.statusCode = http.StatusOK
+	}
+	res, err := lr.w.Write(dat)
+	lr.contentLength += res
+	return res, err
+}
+
+func (lr *loggedResponseWriter) WriteHeader(statusCode int) {
+	lr.statusCode = statusCode
+	lr.w.WriteHeader(statusCode)
 }
 
 type LoggerMiddleware struct {
 	l *logrus.Logger
+	h http.Handler
 }
 
 func (lm *LoggerMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	responseInfo, ok := r.Context().Value(helpers.ResponseInfoKey).(helpers.ResponseInfo)
-	if !ok {
-		lm.l.Errorf("Failed to read response info from context, got %v", r.Context().Value("response_info"))
-		responseInfo = defaultResponseInfo
+	recorder := loggedResponseWriter{
+		contentLength: -1,
+		statusCode:    -1,
+		w:             w,
 	}
+	lm.h.ServeHTTP(&recorder, r)
+
 	method := r.Method
 	route := r.RequestURI
 	currentTime := time.Now().UTC()
@@ -35,5 +56,5 @@ func (lm *LoggerMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		clientIP = "unknown"
 	}
 
-	lm.l.Infof(`%s - [%s] "%s %s %s" %d %d`, clientIP, currentTime.Format(time.RFC3339), method, route, httpVersion, responseInfo.ReturnCode, responseInfo.ContentLength)
+	lm.l.Infof(`%s - [%s] "%s %s %s" %d %d`, clientIP, currentTime.Format(time.RFC3339), method, route, httpVersion, recorder.statusCode, recorder.contentLength)
 }

polyculeconnect/middlewares/middlewarechain.go

@@ -6,18 +6,9 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-type MiddlewareChains struct {
-	handlers []http.Handler
-}
-
-func (mc *MiddlewareChains) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
-	for _, h := range mc.handlers {
-		h.ServeHTTP(rw, r)
-	}
-}
-
-func WithLogger(handler http.Handler, l *logrus.Logger) *MiddlewareChains {
-	return &MiddlewareChains{
-		handlers: []http.Handler{handler, &LoggerMiddleware{l}},
+func WithLogger(handler http.Handler, l *logrus.Logger) http.Handler {
+	return &LoggerMiddleware{
+		l: l,
+		h: handler,
 	}
 }