From 898c8a408159d9114760b4c2265fc54ffe93095d Mon Sep 17 00:00:00 2001
From: vicanso <tree.xie@outlook.com>
Date: Wed, 9 Aug 2023 20:07:02 +0800
Subject: [PATCH] feat: support before response handler

---
 middleware/static_serve.go | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/middleware/static_serve.go b/middleware/static_serve.go
index b55a689..75e13d0 100644
--- a/middleware/static_serve.go
+++ b/middleware/static_serve.go
@@ -74,7 +74,9 @@ type (
 		NotFoundNext bool
 		// 符合该正则则设置为no cache
 		NoCacheRegexp *regexp.Regexp
-		Skipper       elton.Skipper
+		// 响应前的处理(只针对读取到buffer的文件)
+		BeforeResponse func(string, []byte) ([]byte, error)
+		Skipper        elton.Skipper
 	}
 	// FS file system
 	FS struct {
@@ -265,7 +267,7 @@ func NewStaticServe(staticFile StaticFile, config StaticServeConfig) elton.Handl
 		}
 
 		for k, v := range config.Header {
-			c.SetHeader(k, v)
+			c.AddHeader(k, v)
 		}
 		// 未设置cache control
 		// 或文件符合正则
@@ -275,6 +277,14 @@ func NewStaticServe(staticFile StaticFile, config StaticServeConfig) elton.Handl
 		} else {
 			c.SetHeader(elton.HeaderCacheControl, cacheControl)
 		}
+		// 如果有设置before response
+		if config.BeforeResponse != nil && fileBuf != nil {
+			buf, err := config.BeforeResponse(file, fileBuf)
+			if err != nil {
+				return err
+			}
+			fileBuf = buf
+		}
 		if fileBuf != nil {
 			c.StatusCode = http.StatusOK
 			c.BodyBuffer = bytes.NewBuffer(fileBuf)