From cc3b31b0d809b40902f4a6a41e3490a64b6c5e28 Mon Sep 17 00:00:00 2001
From: vicanso <vicansocanbico@gmail.com>
Date: Fri, 27 Nov 2020 21:12:19 +0800
Subject: [PATCH] refactor: logger support to get value from context

---
 middleware/logger.go      |  8 ++++++++
 middleware/logger_test.go | 17 +++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/middleware/logger.go b/middleware/logger.go
index 5037220..07b68c3 100644
--- a/middleware/logger.go
+++ b/middleware/logger.go
@@ -61,6 +61,7 @@ const (
 	payloadSizeHuman = "payload-size-human"
 	requestHeader    = "requestHeader"
 	responseHeader   = "responseHeader"
+	context          = "context"
 	httpProto        = "HTTP"
 	httpsProto       = "HTTPS"
 
@@ -167,6 +168,11 @@ func parseLoggerTags(desc []byte) []*LoggerTag {
 				category: responseHeader,
 				data:     byteSliceToString(k[1:]),
 			})
+		case byte(':'):
+			arr = append(arr, &LoggerTag{
+				category: context,
+				data:     byteSliceToString(k[1:]),
+			})
 		default:
 			arr = append(arr, &LoggerTag{
 				category: byteSliceToString(k),
@@ -223,6 +229,8 @@ func formatLog(c *elton.Context, tags []*LoggerTag, startedAt time.Time) string
 			return c.Request.Header.Get(tag.data)
 		case responseHeader:
 			return c.GetHeader(tag.data)
+		case context:
+			return c.GetString(tag.data)
 		case referer:
 			return c.Request.Referer()
 		case userAgent:
diff --git a/middleware/logger_test.go b/middleware/logger_test.go
index 45cc701..a5cdef4 100644
--- a/middleware/logger_test.go
+++ b/middleware/logger_test.go
@@ -149,6 +149,23 @@ func TestLogger(t *testing.T) {
 		assert.Nil(err)
 	})
 
+	t.Run("context", func(t *testing.T) {
+		config := LoggerConfig{
+			Format: "{:id}",
+			OnLog: func(log string, _ *elton.Context) {
+				assert.Equal("1", log, "get context value fail")
+			},
+		}
+		m := NewLogger(config)
+		c := elton.NewContext(nil, nil)
+		c.Set("id", "1")
+		c.Next = func() error {
+			return nil
+		}
+		err := m(c)
+		assert.Nil(err)
+	})
+
 	t.Run("get log function", func(t *testing.T) {
 		layout := "{host} {remote} {real-ip} {method} {path} {proto} {query} {scheme} {uri} {referer} {userAgent} {size} {size-human} {status} {payload-size} {payload-size-human}"
 		fn := GenerateLog(layout)