From 02ac45383cf9d10d8d75bcfa58037007d39477a9 Mon Sep 17 00:00:00 2001 From: Skyenought <1808644906@qq.com> Date: Tue, 15 Nov 2022 23:47:40 +0800 Subject: [PATCH] :sparkles: Add customTags in logger middleware Config --- middleware/logger/README.md | 18 +++- middleware/logger/config.go | 33 ++++--- middleware/logger/logger.go | 146 +++++++++---------------------- middleware/logger/logger_test.go | 29 ++++++ middleware/logger/tag.go | 132 ++++++++++++++++++++++++++++ 5 files changed, 243 insertions(+), 115 deletions(-) create mode 100644 middleware/logger/tag.go diff --git a/middleware/logger/README.md b/middleware/logger/README.md index 63c1530fda5..7b69d916f9c 100644 --- a/middleware/logger/README.md +++ b/middleware/logger/README.md @@ -11,6 +11,7 @@ Logger middleware for [Fiber](https://github.com/gofiber/fiber) that logs HTTP r - [Logging Request ID](#logging-request-id) - [Changing TimeZone & TimeFormat](#changing-timezone--timeformat) - [Custom File Writer](#custom-file-writer) + - [Add Custom Tags](#add-custom-tags) - [Config](#config) - [Default Config](#default-config-1) - [Constants](#constants) @@ -75,6 +76,16 @@ app.Use(logger.New(logger.Config{ Output: file, })) ``` +### Add Custom Tags +```go +app.Use(logger.New(logger.Config{ + CustomTags: map[string]logger.LogFunc{ + "custom_tag": func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString("it is a custom tag") + }, + }, +})) +``` ## Config ```go @@ -85,11 +96,16 @@ type Config struct { // Optional. Default: nil Next func(c *fiber.Ctx) bool + // CustomTags defines the custom tag action + // + // Optional. Default: map[string]LogFunc{} + CustomTags map[string]LogFunc + // Format defines the logging tags // // Optional. Default: [${time}] ${status} - ${latency} ${method} ${path}\n Format string - + // TimeFormat https://programming.guide/go/format-parse-string-time-date-example.html // // Optional. Default: 15:04:05 diff --git a/middleware/logger/config.go b/middleware/logger/config.go index 0e8aaa5caff..6e112d4ffad 100644 --- a/middleware/logger/config.go +++ b/middleware/logger/config.go @@ -5,7 +5,7 @@ import ( "os" "strings" "time" - + "github.com/gofiber/fiber/v2" ) @@ -15,38 +15,43 @@ type Config struct { // // Optional. Default: nil Next func(c *fiber.Ctx) bool - + // Done is a function that is called after the log string for a request is written to Output, // and pass the log string as parameter. // // Optional. Default: a function that does nothing. Done func(c *fiber.Ctx, logString []byte) - + + // CustomTags defines the custom tag action + // + // Optional. Default: map[string]logFunc + CustomTags map[string]logFunc + // Format defines the logging tags // // Optional. Default: [${time}] ${status} - ${latency} ${method} ${path}\n Format string - + // TimeFormat https://programming.guide/go/format-parse-string-time-date-example.html // // Optional. Default: 15:04:05 TimeFormat string - + // TimeZone can be specified, such as "UTC" and "America/New_York" and "Asia/Chongqing", etc // // Optional. Default: "Local" TimeZone string - + // TimeInterval is the delay before the timestamp is updated // // Optional. Default: 500 * time.Millisecond TimeInterval time.Duration - + // Output is a writer where logs are written // // Default: os.Stdout Output io.Writer - + enableColors bool enableLatency bool timeZoneLocation *time.Location @@ -84,15 +89,15 @@ func configDefault(config ...Config) Config { if len(config) < 1 { return ConfigDefault } - + // Override default config cfg := config[0] - + // Enable colors if no custom format or output is given if validCustomFormat(cfg.Format) && cfg.Output == nil { cfg.enableColors = true } - + // Set default values if cfg.Next == nil { cfg.Next = ConfigDefault.Next @@ -115,5 +120,11 @@ func configDefault(config ...Config) Config { if cfg.Output == nil { cfg.Output = ConfigDefault.Output } + // Set custom tags + if len(cfg.CustomTags) != 0 { + for k, v := range cfg.CustomTags { + tagMap[k] = v + } + } return cfg } diff --git a/middleware/logger/logger.go b/middleware/logger/logger.go index dc348bfebe6..ae3c5aabfdf 100644 --- a/middleware/logger/logger.go +++ b/middleware/logger/logger.go @@ -78,6 +78,24 @@ func New(config ...Config) fiber.Handler { // Check if format contains latency cfg.enableLatency = strings.Contains(cfg.Format, "${latency}") + // Add TagStatus Config + tagMap[TagStatus] = func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + colors := c.App().Config().ColorScheme + if cfg.enableColors { + return buf.WriteString(fmt.Sprintf("%s %3d %s", statusColor(c.Response().StatusCode(), colors), c.Response().StatusCode(), colors.Reset)) + } + return appendInt(buf, c.Response().StatusCode()) + } + + // Add TagMethod Config + tagMap[TagMethod] = func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + colors := c.App().Config().ColorScheme + if cfg.enableColors { + return buf.WriteString(fmt.Sprintf("%s %-7s %s", methodColor(c.Method(), colors), c.Method(), colors.Reset)) + } + return buf.WriteString(c.Method()) + } + // Create template parser tmpl := fasttemplate.New(cfg.Format, "${", "}") @@ -98,6 +116,16 @@ func New(config ...Config) fiber.Handler { // Set PID once pid := strconv.Itoa(os.Getpid()) + tagMap[TagPid] = func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(pid) + } + + tagMap[TagLatency] = func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + ctx := c.Context() + start, _ := time.ParseInLocation(cfg.TimeFormat, ctx.UserValue("start").(string), cfg.timeZoneLocation) + stop, _ := time.ParseInLocation(cfg.TimeFormat, ctx.UserValue("stop").(string), cfg.timeZoneLocation) + return buf.WriteString(fmt.Sprintf("%7v", stop.Sub(start))) + } // Set variables var ( once sync.Once @@ -114,6 +142,11 @@ func New(config ...Config) fiber.Handler { } errPadding := 15 errPaddingStr := strconv.Itoa(errPadding) + + tagMap[TagTime] = func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(timestamp.Load().(string)) + } + // Return new handler return func(c *fiber.Ctx) (err error) { // Don't execute middleware if Next returns true @@ -145,6 +178,7 @@ func New(config ...Config) fiber.Handler { // Set latency start time if cfg.enableLatency { start = time.Now() + c.Context().SetUserValue("start", start.String()) } // Handle request, store err for logging @@ -152,6 +186,7 @@ func New(config ...Config) fiber.Handler { // Manually call error handler if chainErr != nil { + c.Context().SetUserValue("loggerChainError", chainErr.Error()) if err := errHandler(c, chainErr); err != nil { _ = c.SendStatus(fiber.StatusInternalServerError) } @@ -160,6 +195,7 @@ func New(config ...Config) fiber.Handler { // Set latency stop time if cfg.enableLatency { stop = time.Now() + c.Context().SetUserValue("stop", stop.String()) } // Get new buffer @@ -198,109 +234,13 @@ func New(config ...Config) fiber.Handler { // Loop over template tags to replace it with the correct value _, err = tmpl.ExecuteFunc(buf, func(w io.Writer, tag string) (int, error) { - switch tag { - case TagTime: - return buf.WriteString(timestamp.Load().(string)) - case TagReferer: - return buf.WriteString(c.Get(fiber.HeaderReferer)) - case TagProtocol: - return buf.WriteString(c.Protocol()) - case TagPid: - return buf.WriteString(pid) - case TagPort: - return buf.WriteString(c.Port()) - case TagIP: - return buf.WriteString(c.IP()) - case TagIPs: - return buf.WriteString(c.Get(fiber.HeaderXForwardedFor)) - case TagHost: - return buf.WriteString(c.Hostname()) - case TagPath: - return buf.WriteString(c.Path()) - case TagURL: - return buf.WriteString(c.OriginalURL()) - case TagUA: - return buf.WriteString(c.Get(fiber.HeaderUserAgent)) - case TagLatency: - return buf.WriteString(fmt.Sprintf("%7v", stop.Sub(start).Round(time.Millisecond))) - case TagBody: - return buf.Write(c.Body()) - case TagBytesReceived: - return appendInt(buf, len(c.Request().Body())) - case TagBytesSent: - return appendInt(buf, len(c.Response().Body())) - case TagRoute: - return buf.WriteString(c.Route().Path) - case TagStatus: - if cfg.enableColors { - return buf.WriteString(fmt.Sprintf("%s %3d %s", statusColor(c.Response().StatusCode(), colors), c.Response().StatusCode(), colors.Reset)) - } - return appendInt(buf, c.Response().StatusCode()) - case TagResBody: - return buf.Write(c.Response().Body()) - case TagReqHeaders: - reqHeaders := make([]string, 0) - for k, v := range c.GetReqHeaders() { - reqHeaders = append(reqHeaders, k+"="+v) - } - return buf.Write([]byte(strings.Join(reqHeaders, "&"))) - case TagQueryStringParams: - return buf.WriteString(c.Request().URI().QueryArgs().String()) - case TagMethod: - if cfg.enableColors { - return buf.WriteString(fmt.Sprintf("%s %-7s %s", methodColor(c.Method(), colors), c.Method(), colors.Reset)) - } - return buf.WriteString(c.Method()) - case TagBlack: - return buf.WriteString(colors.Black) - case TagRed: - return buf.WriteString(colors.Red) - case TagGreen: - return buf.WriteString(colors.Green) - case TagYellow: - return buf.WriteString(colors.Yellow) - case TagBlue: - return buf.WriteString(colors.Blue) - case TagMagenta: - return buf.WriteString(colors.Magenta) - case TagCyan: - return buf.WriteString(colors.Cyan) - case TagWhite: - return buf.WriteString(colors.White) - case TagReset: - return buf.WriteString(colors.Reset) - case TagError: - if chainErr != nil { - return buf.WriteString(chainErr.Error()) - } - return buf.WriteString("-") - default: - // Check if we have a value tag i.e.: "reqHeader:x-key" - switch { - case strings.HasPrefix(tag, TagReqHeader): - return buf.WriteString(c.Get(tag[10:])) - case strings.HasPrefix(tag, TagHeader): - return buf.WriteString(c.Get(tag[7:])) - case strings.HasPrefix(tag, TagRespHeader): - return buf.WriteString(c.GetRespHeader(tag[11:])) - case strings.HasPrefix(tag, TagQuery): - return buf.WriteString(c.Query(tag[6:])) - case strings.HasPrefix(tag, TagForm): - return buf.WriteString(c.FormValue(tag[5:])) - case strings.HasPrefix(tag, TagCookie): - return buf.WriteString(c.Cookies(tag[7:])) - case strings.HasPrefix(tag, TagLocals): - switch v := c.Locals(tag[7:]).(type) { - case []byte: - return buf.Write(v) - case string: - return buf.WriteString(v) - case nil: - return 0, nil - default: - return buf.WriteString(fmt.Sprintf("%v", v)) - } - } + if logFunc, ok := tagMap[tag]; ok { + return logFunc(buf, c, w, tag) + } + split := strings.SplitAfter(tag, ":") + if ok := strings.Contains(tag, split[0]); ok && len(split) > 1 { + logFunc, _ := tagMap[split[0]] + return logFunc(buf, c, w, tag) } return 0, nil }) diff --git a/middleware/logger/logger_test.go b/middleware/logger/logger_test.go index ba566d2c2ea..e313a9173ea 100644 --- a/middleware/logger/logger_test.go +++ b/middleware/logger/logger_test.go @@ -383,3 +383,32 @@ func Test_ReqHeader_Header(t *testing.T) { utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode) utils.AssertEqual(t, "Hello fiber!", buf.String()) } + +// go test -run Test_CustomTags +func Test_CustomTags(t *testing.T) { + customTag := "it is a custom tag" + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app := fiber.New() + app.Use(New(Config{ + Format: "${custom_tag}", + CustomTags: map[string]logFunc{ + "custom_tag": func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(customTag) + }, + }, + Output: buf, + })) + app.Get("/", func(c *fiber.Ctx) error { + return c.SendString("Hello fiber!") + }) + reqHeaderReq := httptest.NewRequest("GET", "/", nil) + reqHeaderReq.Header.Add("test", "Hello fiber!") + resp, err := app.Test(reqHeaderReq) + + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode) + utils.AssertEqual(t, customTag, buf.String()) +} diff --git a/middleware/logger/tag.go b/middleware/logger/tag.go new file mode 100644 index 00000000000..b6c2a82d8c8 --- /dev/null +++ b/middleware/logger/tag.go @@ -0,0 +1,132 @@ +package logger + +import ( + "fmt" + "io" + "strings" + + "github.com/gofiber/fiber/v2" + "github.com/gofiber/fiber/v2/internal/bytebufferpool" +) + +type logFunc func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) + +// tagMap use to storage how to log content when parse Tag. +var tagMap = map[string]logFunc{ + TagReferer: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Get(fiber.HeaderReferer)) + }, + TagProtocol: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Protocol()) + }, + TagPort: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Port()) + }, + TagIP: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.IP()) + }, + TagIPs: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Get(fiber.HeaderXForwardedFor)) + }, + TagHost: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Hostname()) + }, + TagPath: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Path()) + }, + TagURL: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.OriginalURL()) + }, + TagUA: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Get(fiber.HeaderUserAgent)) + }, + TagBody: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.Write(c.Body()) + }, + TagBytesReceived: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return appendInt(buf, len(c.Request().Body())) + }, + TagBytesSent: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return appendInt(buf, len(c.Response().Body())) + }, + TagRoute: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Route().Path) + }, + TagResBody: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.Write(c.Response().Body()) + }, + TagReqHeaders: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + reqHeaders := make([]string, 0) + for k, v := range c.GetReqHeaders() { + reqHeaders = append(reqHeaders, k+"="+v) + } + return buf.Write([]byte(strings.Join(reqHeaders, "&"))) + }, + TagQueryStringParams: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Request().URI().QueryArgs().String()) + }, + + TagBlack: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Black) + }, + TagRed: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Red) + }, + TagGreen: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Green) + }, + TagYellow: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Yellow) + }, + TagBlue: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Blue) + }, + TagMagenta: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Magenta) + }, + TagCyan: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Cyan) + }, + TagWhite: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.White) + }, + TagReset: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.App().Config().ColorScheme.Reset) + }, + TagError: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + if s := c.Context().UserValue("loggerChainError").(string); s != "" { + return buf.WriteString(s) + } + return buf.WriteString("-") + }, + TagReqHeader: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Get(tag[10:])) + }, + TagHeader: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Get(tag[7:])) + }, + TagRespHeader: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.GetRespHeader(tag[11:])) + }, + TagQuery: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Query(tag[6:])) + }, + TagForm: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.FormValue(tag[5:])) + }, + TagCookie: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + return buf.WriteString(c.Cookies(tag[7:])) + }, + TagLocals: func(buf *bytebufferpool.ByteBuffer, c *fiber.Ctx, w io.Writer, tag string) (int, error) { + switch v := c.Locals(tag[7:]).(type) { + case []byte: + return buf.Write(v) + case string: + return buf.WriteString(v) + case nil: + return 0, nil + default: + return buf.WriteString(fmt.Sprintf("%v", v)) + } + }, +}