add CacheControl to Static config

gofiber · Oct 5, 2022 · 5a67c2a · 5a67c2a
1 parent e85a29a
commit 5a67c2a
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 0 deletions.
diff --git a/app.go b/app.go
@@ -421,6 +421,11 @@ type Static struct {
 	// Optional. Default value 0.
 	MaxAge int `json:"max_age"`
 
+	// SetHeaders defines a function that sets custom headers on response.
+	//
+	// Optional. Default: nil
+	SetHeaders func(c *Ctx)
+
 	// Next defines a function to skip this middleware when returned true.
 	//
 	// Optional. Default: nil

diff --git a/app_test.go b/app_test.go
@@ -767,6 +767,31 @@ func Test_App_Static_MaxAge(t *testing.T) {
 	utils.AssertEqual(t, "public, max-age=100", resp.Header.Get(HeaderCacheControl), "CacheControl Control")
 }
 
+// go test -run Test_App_Static_Custom_CacheControl
+func Test_App_Static_Custom_CacheControl(t *testing.T) {
+	app := New()
+
+	app.Static("/", "./.github", Static{SetHeaders: func(c *Ctx) {
+		if strings.Contains(string(c.Response().Header.ContentType()), "text/html") {
+			c.Response().Header.Set("Cache-Control", "no-cache, no-store, must-revalidate")
+		}
+	}})
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/index.html", nil))
+	utils.AssertEqual(t, nil, err, "app.Test(req)")
+	utils.AssertEqual(t, 200, resp.StatusCode, "Status code")
+	utils.AssertEqual(t, false, resp.Header.Get(HeaderContentLength) == "")
+	utils.AssertEqual(t, "text/html; charset=utf-8", resp.Header.Get(HeaderContentType))
+	utils.AssertEqual(t, "no-cache, no-store, must-revalidate", resp.Header.Get(HeaderCacheControl), "CacheControl Control")
+
+	normal_resp, normal_err := app.Test(httptest.NewRequest("GET", "/config.yml", nil))
+	utils.AssertEqual(t, nil, normal_err, "app.Test(req)")
+	utils.AssertEqual(t, 200, normal_resp.StatusCode, "Status code")
+	utils.AssertEqual(t, false, normal_resp.Header.Get(HeaderContentLength) == "")
+	utils.AssertEqual(t, "text/plain; charset=utf-8", normal_resp.Header.Get(HeaderContentType))
+	utils.AssertEqual(t, "", normal_resp.Header.Get(HeaderCacheControl), "CacheControl Control")
+}
+
 // go test -run Test_App_Static_Download
 func Test_App_Static_Download(t *testing.T) {
 	app := New()

diff --git a/router.go b/router.go
@@ -357,6 +357,7 @@ func (app *App) registerStatic(prefix, root string, config ...Static) Router {
 
 	// Set config if provided
 	var cacheControlValue string
+	var setHeaders func(c *Ctx)
 	if len(config) > 0 {
 		maxAge := config[0].MaxAge
 		if maxAge > 0 {
@@ -369,6 +370,7 @@ func (app *App) registerStatic(prefix, root string, config ...Static) Router {
 		if config[0].Index != "" {
 			fs.IndexNames = []string{config[0].Index}
 		}
+		setHeaders = config[0].SetHeaders
 	}
 	fileHandler := fs.NewRequestHandler()
 	handler := func(c *Ctx) error {
@@ -388,6 +390,9 @@ func (app *App) registerStatic(prefix, root string, config ...Static) Router {
 			if len(cacheControlValue) > 0 {
 				c.fasthttp.Response.Header.Set(HeaderCacheControl, cacheControlValue)
 			}
+			if setHeaders != nil {
+				setHeaders(c)
+			}
 			return nil
 		}
 		// Reset response to default