Better middleware

Author: kalverra

parrot/errors.go

@@ -6,13 +6,16 @@ import (
 )
 
 var (
-	ErrNilRoute         = errors.New("route is nil")
-	ErrNoMethod         = errors.New("no method specified")
-	ErrInvalidPath      = errors.New("invalid path")
-	ErrNoResponse       = errors.New("route must have a handler or some response")
-	ErrOnlyOneResponse  = errors.New("route can only have one response type")
-	ErrResponseMarshal  = errors.New("unable to marshal response body to JSON")
-	ErrRouteNotFound    = errors.New("route not found")
+	ErrNilRoute        = errors.New("route is nil")
+	ErrNoMethod        = errors.New("no method specified")
+	ErrInvalidPath     = errors.New("invalid path")
+	ErrNoResponse      = errors.New("route must have a handler or some response")
+	ErrOnlyOneResponse = errors.New("route can only have one response type")
+	ErrResponseMarshal = errors.New("unable to marshal response body to JSON")
+	ErrRouteNotFound   = errors.New("route not found")
+
+	ErrNoRecorderURL    = errors.New("no recorder URL specified")
+	ErrNilRecorder      = errors.New("recorder is nil")
 	ErrRecorderNotFound = errors.New("recorder not found")
 )
 

parrot/go.mod

@@ -4,6 +4,7 @@ go 1.23.4
 
 require (
 	github.com/go-resty/resty/v2 v2.16.3
+	github.com/google/uuid v1.6.0
 	github.com/rs/zerolog v1.33.0
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
@@ -15,6 +16,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/rs/xid v1.5.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/net v0.33.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect

parrot/go.sum

@@ -5,6 +5,8 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/go-resty/resty/v2 v2.16.3 h1:zacNT7lt4b8M/io2Ahj6yPypL7bqx9n1iprfQuodV+E=
 github.com/go-resty/resty/v2 v2.16.3/go.mod h1:hkJtXbA2iKHzJheXYvQ8snQES5ZLGKMwQ07xAwp/fiA=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
@@ -15,6 +17,7 @@ github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rs/xid v1.5.0 h1:mKX4bl4iPYJtEIxp6CYiUuLQ/8DYMoz0PUdtGgMFRVc=
 github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
 github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
 github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=

parrot/parrot.go

@@ -18,7 +18,9 @@ import (
 	"time"
 
 	"github.com/go-resty/resty/v2"
+	"github.com/google/uuid"
 	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/hlog"
 )
 
 // Route holds information about the mock route configuration
@@ -173,7 +175,7 @@ func Wake(options ...ServerOption) (*Server, error) {
 		if p.jsonLogs {
 			writers = append(writers, os.Stderr)
 		} else {
-			consoleOut := zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339Nano}
+			consoleOut := zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: "2006-01-02T15:04:05.000"}
 			writers = append(writers, consoleOut)
 		}
 
@@ -208,7 +210,7 @@ func Wake(options ...ServerOption) (*Server, error) {
 	p.server = &http.Server{
 		ReadHeaderTimeout: 5 * time.Second,
 		Addr:              listener.Addr().String(),
-		Handler:           mux,
+		Handler:           p.loggingMiddleware(mux),
 	}
 
 	if err = p.load(); err != nil {
@@ -290,69 +292,71 @@ func (p *Server) Register(route *Route) error {
 	p.routesMu.Lock()
 	defer p.routesMu.Unlock()
 	p.routes[route.ID()] = route
+	p.log.Info().
+		Str("Route ID", route.ID()).
+		Str("Path", route.Path).
+		Str("Method", route.Method).
+		Msg("Route registered")
 
 	return nil
 }
 
 // registerRouteHandler handles the dynamic route registration.
 func (p *Server) registerRouteHandler(w http.ResponseWriter, r *http.Request) {
-	const parrotPath = "/register"
+	registerLogger := zerolog.Ctx(r.Context())
 	if r.Method == http.MethodDelete {
 		var routeRequest *RouteRequest
 		if err := json.NewDecoder(r.Body).Decode(&routeRequest); err != nil {
 			http.Error(w, "Invalid request body", http.StatusBadRequest)
+			registerLogger.Debug().Err(err).Msg("Failed to decode request body")
 			return
 		}
 		defer r.Body.Close()
 
 		if routeRequest.ID == "" {
-			err := errors.New("ID required")
-			http.Error(w, err.Error(), http.StatusBadRequest)
+			http.Error(w, "Route ID required", http.StatusBadRequest)
+			registerLogger.Debug().Msg("No Route ID provided")
 			return
 		}
 
 		err := p.Unregister(routeRequest.ID)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
-			p.log.Trace().Err(err).Str("Path", parrotPath).Msg("Failed to unregister route")
+			registerLogger.Debug().Err(err).Msg("Failed to unregister route")
 			return
 		}
 
 		w.WriteHeader(http.StatusNoContent)
-		p.log.Info().
+		registerLogger.Info().
 			Str("Route ID", routeRequest.ID).
-			Str("Parrot Path", parrotPath).
 			Msg("Route unregistered")
 	} else if r.Method == http.MethodPost {
 		var route *Route
 		if err := json.NewDecoder(r.Body).Decode(&route); err != nil {
 			http.Error(w, "Invalid request body", http.StatusBadRequest)
+			registerLogger.Debug().Err(err).Msg("Failed to decode request body")
 			return
 		}
 		defer r.Body.Close()
 
 		if route.Method == "" || route.Path == "" {
 			err := errors.New("Method and path are required")
 			http.Error(w, err.Error(), http.StatusBadRequest)
+			registerLogger.Debug().Err(err).Msg("Method and path are required")
 			return
 		}
 
 		err := p.Register(route)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
-			p.log.Trace().Err(err).Msg("Failed to register route")
+			registerLogger.Debug().Err(err).Msg("Failed to register route")
 			return
 		}
 
 		w.WriteHeader(http.StatusCreated)
-		p.log.Info().
-			Str("Parrot Path", parrotPath).
-			Str("Route Path", route.Path).
-			Str("Method", route.Method).
-			Msg("Route registered")
 	} else {
 		http.Error(w, "Invalid method, only use POST or DELETE", http.StatusMethodNotAllowed)
-		p.log.Trace().Str("Method", r.Method).Msg("Invalid method")
+		registerLogger.Debug().Msg("Invalid method")
 		return
 	}
 }
@@ -361,34 +365,44 @@ func (p *Server) registerRouteHandler(w http.ResponseWriter, r *http.Request) {
 func (p *Server) Record(recorder *Recorder) error {
 	p.recordersMu.Lock()
 	defer p.recordersMu.Unlock()
+	if recorder == nil {
+		return ErrNilRecorder
+	}
+	if recorder.URL == "" {
+		return ErrNoRecorderURL
+	}
+	_, err := url.Parse(recorder.URL)
+	if err != nil {
+		return fmt.Errorf("failed to parse recorder URL: %w", err)
+	}
 	p.recorderHooks = append(p.recorderHooks, recorder.URL)
 	return nil
 }
 
 func (p *Server) recordHandler(w http.ResponseWriter, r *http.Request) {
-	const parrotPath = "/record"
+	recordLogger := zerolog.Ctx(r.Context())
 	if r.Method != http.MethodPost {
 		http.Error(w, "Invalid method, only use POST or DELETE", http.StatusMethodNotAllowed)
-		p.log.Trace().Str("Method", r.Method).Msg("Invalid method")
+		recordLogger.Debug().Msg("Invalid method")
 		return
 	}
 
 	var recorder *Recorder
 	if err := json.NewDecoder(r.Body).Decode(&recorder); err != nil {
 		http.Error(w, "Invalid request body", http.StatusBadRequest)
-		p.log.Trace().Err(err).Str("Parrot Path", parrotPath).Msg("Failed to decode request body")
+		recordLogger.Err(err).Msg("Failed to decode request body")
 		return
 	}
 
 	err := p.Record(recorder)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
-		p.log.Trace().Err(err).Str("Parrot Path", parrotPath).Msg("Failed to add recorder")
+		recordLogger.Debug().Err(err).Msg("Failed to add recorder")
 		return
 	}
 
 	w.WriteHeader(http.StatusCreated)
-	p.log.Info().Str("Recorder URL", recorder.URL).Str("Parrot Path", parrotPath).Msg("Recorder added")
+	recordLogger.Info().Str("Recorder URL", recorder.URL).Msg("Recorder added")
 }
 
 // Unregister removes a route from the parrot
@@ -423,78 +437,98 @@ func (p *Server) dynamicHandler(w http.ResponseWriter, r *http.Request) {
 	route, exists := p.routes[r.Method+":"+r.URL.Path]
 	p.routesMu.RUnlock()
 
+	dynamicLogger := zerolog.Ctx(r.Context())
+	if !exists {
+		http.NotFound(w, r)
+		dynamicLogger.Debug().Msg("Route not found")
+		return
+	}
+
+	requestID := uuid.New().String()[0:8]
+	dynamicLogger.UpdateContext(func(c zerolog.Context) zerolog.Context {
+		return c.Str("Request ID", requestID).Str("Route ID", route.ID())
+	})
+
+	requestBody, err := io.ReadAll(r.Body)
+	if err != nil {
+		dynamicLogger.Debug().
+			Err(err).
+			Msg("Failed to read request body")
+		http.Error(w, "Failed to read request body", http.StatusInternalServerError)
+		return
+	}
+
 	routeCall := &RouteCall{
 		RouteID: r.Method + ":" + r.URL.Path,
-		Request: r,
+		Request: &RouteCallRequest{
+			Method: r.Method,
+			URL:    r.URL,
+			Header: r.Header,
+			Body:   requestBody,
+		},
 	}
 	recordingWriter := newResponseWriterRecorder(w)
 
 	defer func() {
-		routeCall.Response = recordingWriter.Result()
+		res := recordingWriter.Result()
+		resBody, err := io.ReadAll(res.Body)
+		if err != nil {
+			dynamicLogger.Debug().Err(err).Msg("Failed to read response body")
+			http.Error(w, "Failed to read response body", http.StatusInternalServerError)
+			return
+		}
+
+		routeCall.Response = &RouteCallResponse{
+			StatusCode: res.StatusCode,
+			Header:     res.Header,
+			Body:       resBody,
+		}
 		p.sendToRecorders(routeCall)
 	}()
 
-	if !exists { // Route not found
-		http.NotFound(recordingWriter, r)
-		p.log.Trace().Str("Remote Addr", r.RemoteAddr).Str("Path", r.URL.Path).Str("Method", r.Method).Msg("Route not found")
-		return
-	}
-
 	// Let the custom handler take over if it exists
 	if route.Handler != nil {
-		p.log.Trace().Str("Remote Addr", r.RemoteAddr).Str("Path", r.URL.Path).Str("Method", r.Method).Msg("Calling route handler")
+		dynamicLogger.Debug().Msg("Calling route handler")
 		route.Handler(recordingWriter, r)
 		return
 	}
 
-	recordingWriter.WriteHeader(route.ResponseStatusCode)
-
 	if route.RawResponseBody != "" {
 		if _, err := w.Write([]byte(route.RawResponseBody)); err != nil {
-			p.log.Trace().Err(err).Str("Remote Addr", r.RemoteAddr).Str("Path", r.URL.Path).Str("Method", r.Method).Msg("Failed to write response")
+			dynamicLogger.Debug().Err(err).Msg("Failed to write response")
 			http.Error(recordingWriter, "Failed to write response", http.StatusInternalServerError)
 			return
 		}
-		p.log.Trace().
-			Str("Remote Addr", r.RemoteAddr).
+		dynamicLogger.Debug().
 			Str("Response", route.RawResponseBody).
-			Str("Path", r.URL.Path).
-			Str("Method", r.Method).
 			Msg("Returned raw response")
+		recordingWriter.WriteHeader(route.ResponseStatusCode)
 		return
 	}
 
 	if route.ResponseBody != nil {
 		rawJSON, err := json.Marshal(route.ResponseBody)
 		if err != nil {
-			p.log.Trace().Err(err).
-				Str("Remote Addr", r.RemoteAddr).
-				Str("Path", r.URL.Path).
-				Str("Method", r.Method).
-				Msg("Failed to marshal JSON response")
+			dynamicLogger.Debug().Err(err).Msg("Failed to marshal JSON response")
 			http.Error(recordingWriter, "Failed to marshal response into json", http.StatusInternalServerError)
 			return
 		}
 		if _, err = w.Write(rawJSON); err != nil {
-			p.log.Trace().Err(err).
+			dynamicLogger.Debug().Err(err).
 				RawJSON("Response", rawJSON).
-				Str("Remote Addr", r.RemoteAddr).
-				Str("Path", r.URL.Path).
-				Str("Method", r.Method).
 				Msg("Failed to write response")
 			http.Error(recordingWriter, "Failed to write JSON response", http.StatusInternalServerError)
 			return
 		}
-		p.log.Trace().
-			Str("Remote Addr", r.RemoteAddr).
+		dynamicLogger.Debug().
 			RawJSON("Response", rawJSON).
-			Str("Path", r.URL.Path).
-			Str("Method", r.Method).
 			Msg("Returned JSON response")
+		recordingWriter.WriteHeader(route.ResponseStatusCode)
 		return
 	}
 
-	p.log.Error().Str("Remote Addr", r.RemoteAddr).Str("Path", r.URL.Path).Str("Method", r.Method).Msg("Route has no response")
+	dynamicLogger.Error().Msg("Route has no response")
+	http.Error(recordingWriter, "Route has no response", http.StatusInternalServerError)
 }
 
 // load loads all registered routes from a file.
@@ -554,12 +588,15 @@ func (p *Server) save() error {
 func (p *Server) sendToRecorders(routeCall *RouteCall) {
 	p.recordersMu.RLock()
 	defer p.recordersMu.RUnlock()
+	if len(p.recorderHooks) == 0 {
+		return
+	}
 
 	client := resty.New()
+	p.log.Trace().Strs("Recorders", p.recorderHooks).Str("Route ID", routeCall.RouteID).Msg("Sending route call to recorders")
 
 	for _, hook := range p.recorderHooks {
 		go func(hook string) {
-			p.log.Trace().Str("Recorder Hook", hook).Msg("Sending route call to recorder")
 			resp, err := client.R().SetBody(routeCall).Post(hook)
 			if err != nil {
 				p.log.Error().Err(err).Str("Recorder Hook", hook).Msg("Failed to send route call to recorder")
@@ -573,11 +610,30 @@ func (p *Server) sendToRecorders(routeCall *RouteCall) {
 					Msg("Failed to send route call to recorder")
 				return
 			}
-			p.log.Debug().Str("Recorder Hook", hook).Msg("Route call sent to recorder")
+			p.log.Trace().Str("Route ID", routeCall.RouteID).Str("Recorder Hook", hook).Msg("Route call sent to recorder")
 		}(hook)
 	}
 }
 
+func (p *Server) loggingMiddleware(next http.Handler) http.Handler {
+	h := hlog.NewHandler(p.log)
+
+	accessHandler := hlog.AccessHandler(
+		func(r *http.Request, status, size int, duration time.Duration) {
+			hlog.FromRequest(r).Trace().
+				Str("Method", r.Method).
+				Stringer("URL", r.URL).
+				Int("Status Code", status).
+				Int("Response Size Bytes", size).
+				Str("Duration", duration.String()).
+				Str("Remote Addr", r.RemoteAddr).
+				Msg("Handled request")
+		},
+	)
+
+	return h(accessHandler(next))
+}
+
 var pathRegex = regexp.MustCompile(`^\/[a-zA-Z0-9\-._~%!$&'()*+,;=:@\/]*$`)
 
 func isValidPath(path string) bool {