Move Logging to Middleware Package (#1070)

* Use a specialized ResponseWriter in middleware

* Track User & Upstream in RequestScope

* Wrap responses in our custom ResponseWriter

* Add tests for logging middleware

* Inject upstream metadata into request scope

* Use custom ResponseWriter only in logging middleware

* Assume RequestScope is never nil

pkg/middleware/request_logger.go

@@ -0,0 +1,110 @@
+package middleware
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+	"time"
+
+	"github.com/justinas/alice"
+	middlewareapi "github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/middleware"
+	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/logger"
+)
+
+// NewRequestLogger returns middleware which logs requests
+// It uses a custom ResponseWriter to track status code & response size details
+func NewRequestLogger() alice.Constructor {
+	return requestLogger
+}
+
+func requestLogger(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		startTime := time.Now()
+		url := *req.URL
+
+		responseLogger := &loggingResponse{ResponseWriter: rw}
+		next.ServeHTTP(responseLogger, req)
+
+		scope := middlewareapi.GetRequestScope(req)
+		// If scope is nil, this will panic.
+		// A scope should always be injected before this handler is called.
+		logger.PrintReq(
+			getUser(scope),
+			scope.Upstream,
+			req,
+			url,
+			startTime,
+			responseLogger.Status(),
+			responseLogger.Size(),
+		)
+	})
+}
+
+func getUser(scope *middlewareapi.RequestScope) string {
+	session := scope.Session
+	if session != nil {
+		if session.Email != "" {
+			return session.Email
+		}
+		return session.User
+	}
+	return ""
+}
+
+// loggingResponse is a custom http.ResponseWriter that allows tracking certain
+// details for request logging.
+type loggingResponse struct {
+	http.ResponseWriter
+
+	status int
+	size   int
+}
+
+// Write writes the response using the ResponseWriter
+func (r *loggingResponse) Write(b []byte) (int, error) {
+	if r.status == 0 {
+		// The status will be StatusOK if WriteHeader has not been called yet
+		r.status = http.StatusOK
+	}
+	size, err := r.ResponseWriter.Write(b)
+	r.size += size
+	return size, err
+}
+
+// WriteHeader writes the status code for the Response
+func (r *loggingResponse) WriteHeader(s int) {
+	r.ResponseWriter.WriteHeader(s)
+	r.status = s
+}
+
+// Hijack implements the `http.Hijacker` interface that actual ResponseWriters
+// implement to support websockets
+func (r *loggingResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
+		return hj.Hijack()
+	}
+	return nil, nil, errors.New("http.Hijacker is not available on writer")
+}
+
+// Flush sends any buffered data to the client. Implements the `http.Flusher`
+// interface
+func (r *loggingResponse) Flush() {
+	if flusher, ok := r.ResponseWriter.(http.Flusher); ok {
+		if r.status == 0 {
+			// The status will be StatusOK if WriteHeader has not been called yet
+			r.status = http.StatusOK
+		}
+		flusher.Flush()
+	}
+}
+
+// Status returns the response status code
+func (r *loggingResponse) Status() int {
+	return r.status
+}
+
+// Size returns the response size
+func (r *loggingResponse) Size() int {
+	return r.size
+}