From fdbfb8b43d6b27233887a254fe932a5d556986b3 Mon Sep 17 00:00:00 2001
From: Cedric Bail <cedric.bail@appdirect.com>
Date: Mon, 8 Jan 2024 10:59:21 -0700
Subject: [PATCH] Link logger with gin context.

---
 logger.go      | 26 +++++++++++++++++++-------
 logger_test.go | 11 +++++++++++
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/logger.go b/logger.go
index dff00a0..f1ec93c 100644
--- a/logger.go
+++ b/logger.go
@@ -32,6 +32,8 @@ type config struct {
 	serverErrorLevel zerolog.Level
 }
 
+const loggerKey = "_gin-contrib/logger_"
+
 var isTerm bool = isatty.IsTerminal(os.Stdout.Fd())
 
 // SetLogger initializes the logging middleware.
@@ -79,9 +81,7 @@ func SetLogger(opts ...Option) gin.HandlerFunc {
 			path = path + "?" + raw
 		}
 
-		c.Next()
 		track := true
-
 		if _, ok := skip[path]; ok {
 			track = false
 		}
@@ -97,6 +97,17 @@ func SetLogger(opts ...Option) gin.HandlerFunc {
 			}
 		}
 
+		if track {
+			l = l.With().
+				Str("method", c.Request.Method).
+				Str("path", path).
+				Str("ip", c.ClientIP()).
+				Str("user_agent", c.Request.UserAgent()).Logger()
+		}
+		c.Set(loggerKey, l)
+
+		c.Next()
+
 		if track {
 			end := time.Now()
 			if cfg.utc {
@@ -106,11 +117,7 @@ func SetLogger(opts ...Option) gin.HandlerFunc {
 
 			l = l.With().
 				Int("status", c.Writer.Status()).
-				Str("method", c.Request.Method).
-				Str("path", path).
-				Str("ip", c.ClientIP()).
-				Dur("latency", latency).
-				Str("user_agent", c.Request.UserAgent()).Logger()
+				Dur("latency", latency).Logger()
 
 			msg := "Request"
 			if len(c.Errors) > 0 {
@@ -141,3 +148,8 @@ func SetLogger(opts ...Option) gin.HandlerFunc {
 func ParseLevel(levelStr string) (zerolog.Level, error) {
 	return zerolog.ParseLevel(levelStr)
 }
+
+// Get return the logger associated with a gin context
+func Get(c *gin.Context) zerolog.Logger {
+	return c.MustGet(loggerKey).(zerolog.Logger)
+}
diff --git a/logger_test.go b/logger_test.go
index 0bf0e35..e869b40 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -67,6 +67,17 @@ func TestLogger(t *testing.T) {
 	assert.Contains(t, buffer.String(), "/example")
 	assert.Contains(t, buffer.String(), "ERR")
 	assert.Contains(t, buffer.String(), "path=/example?a=100")
+
+	buffer.Reset()
+	r.GET("/example-with-additional-log", func(ctx *gin.Context) {
+		l := Get(ctx)
+		l.Info().Msg("additional log")
+	})
+	performRequest(r, "GET", "/example-with-additional-log")
+	assert.Contains(t, buffer.String(), "200")
+	assert.Contains(t, buffer.String(), "GET")
+	assert.Contains(t, buffer.String(), "/example-with-additional-log")
+	assert.Contains(t, buffer.String(), "additional log")
 }
 
 func TestLoggerWithLogger(t *testing.T) {