From 3157fb5f1c45616173da443866e778d8e7c9a3ca Mon Sep 17 00:00:00 2001
From: Andy Pan <panjf2000@gmail.com>
Date: Tue, 15 Nov 2022 19:14:32 +0800
Subject: [PATCH] :sparkles: Add callback function for middleware/logger
 (#2219)

* :sparkles: Add callback function for middleware/logger

* Refine test code
---
 middleware/logger/config.go      | 10 ++++++++++
 middleware/logger/logger.go      | 12 +++++++++---
 middleware/logger/logger_test.go | 25 ++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/middleware/logger/config.go b/middleware/logger/config.go
index 08c5ad51d9..0e8aaa5caf 100644
--- a/middleware/logger/config.go
+++ b/middleware/logger/config.go
@@ -16,6 +16,12 @@ 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)
+
 	// Format defines the logging tags
 	//
 	// Optional. Default: [${time}] ${status} - ${latency} ${method} ${path}\n
@@ -49,6 +55,7 @@ type Config struct {
 // ConfigDefault is the default config
 var ConfigDefault = Config{
 	Next:         nil,
+	Done:         func(c *fiber.Ctx, logString []byte) {},
 	Format:       "[${time}] ${status} - ${latency} ${method} ${path}\n",
 	TimeFormat:   "15:04:05",
 	TimeZone:     "Local",
@@ -90,6 +97,9 @@ func configDefault(config ...Config) Config {
 	if cfg.Next == nil {
 		cfg.Next = ConfigDefault.Next
 	}
+	if cfg.Done == nil {
+		cfg.Done = ConfigDefault.Done
+	}
 	if cfg.Format == "" {
 		cfg.Format = ConfigDefault.Format
 	}
diff --git a/middleware/logger/logger.go b/middleware/logger/logger.go
index 566b685d22..dc348bfebe 100644
--- a/middleware/logger/logger.go
+++ b/middleware/logger/logger.go
@@ -10,12 +10,13 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/gofiber/fiber/v2"
-	"github.com/gofiber/fiber/v2/internal/bytebufferpool"
-	"github.com/gofiber/fiber/v2/internal/fasttemplate"
 	"github.com/mattn/go-colorable"
 	"github.com/mattn/go-isatty"
 	"github.com/valyala/fasthttp"
+
+	"github.com/gofiber/fiber/v2"
+	"github.com/gofiber/fiber/v2/internal/bytebufferpool"
+	"github.com/gofiber/fiber/v2/internal/fasttemplate"
 )
 
 // Logger variables
@@ -186,6 +187,8 @@ func New(config ...Config) fiber.Handler {
 			// Write buffer to output
 			_, _ = cfg.Output.Write(buf.Bytes())
 
+			cfg.Done(c, buf.Bytes())
+
 			// Put buffer back to pool
 			bytebufferpool.Put(buf)
 
@@ -315,6 +318,9 @@ func New(config ...Config) fiber.Handler {
 			}
 		}
 		mu.Unlock()
+
+		cfg.Done(c, buf.Bytes())
+
 		// Put buffer back to pool
 		bytebufferpool.Put(buf)
 
diff --git a/middleware/logger/logger_test.go b/middleware/logger/logger_test.go
index 43522fa300..ba566d2c2e 100644
--- a/middleware/logger/logger_test.go
+++ b/middleware/logger/logger_test.go
@@ -1,6 +1,7 @@
 package logger
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -10,11 +11,12 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/valyala/fasthttp"
+
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/internal/bytebufferpool"
 	"github.com/gofiber/fiber/v2/middleware/requestid"
 	"github.com/gofiber/fiber/v2/utils"
-	"github.com/valyala/fasthttp"
 )
 
 // go test -run Test_Logger
@@ -99,6 +101,27 @@ func Test_Logger_Next(t *testing.T) {
 	utils.AssertEqual(t, fiber.StatusNotFound, resp.StatusCode)
 }
 
+// go test -run Test_Logger_Done
+func Test_Logger_Done(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	app := fiber.New()
+	app.Use(New(Config{
+		Done: func(c *fiber.Ctx, logString []byte) {
+			if c.Response().StatusCode() == fiber.StatusOK {
+				buf.Write(logString)
+			}
+		},
+	})).Get("/logging", func(ctx *fiber.Ctx) error {
+		return ctx.SendStatus(fiber.StatusOK)
+	})
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/logging", nil))
+
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+	utils.AssertEqual(t, true, buf.Len() > 0)
+}
+
 // go test -run Test_Logger_ErrorTimeZone
 func Test_Logger_ErrorTimeZone(t *testing.T) {
 	app := fiber.New()