From 30783dd25d0cd413a611f68aa6a09effd6d2da11 Mon Sep 17 00:00:00 2001
From: kiyon <kiyonlin@163.com>
Date: Thu, 17 Sep 2020 12:56:11 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=B7=20add=20filesystem=20test=20cases?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 middleware/filesystem/filesystem.go      | 54 +++++++++++++-----------
 middleware/filesystem/filesystem_test.go | 39 ++++++++++++++++-
 2 files changed, 66 insertions(+), 27 deletions(-)

diff --git a/middleware/filesystem/filesystem.go b/middleware/filesystem/filesystem.go
index 3c1414d6..458e30fb 100644
--- a/middleware/filesystem/filesystem.go
+++ b/middleware/filesystem/filesystem.go
@@ -47,23 +47,26 @@ var ConfigDefault = Config{
 }
 
 // New creates a new middleware handler
-func New(config Config) fiber.Handler {
-	// Set config
-	cfg := config
+func New(config ...Config) fiber.Handler {
+	// Set default config
+	cfg := ConfigDefault
 
-	// Set default values
-	if cfg.Next == nil {
-		cfg.Next = ConfigDefault.Next
-	}
-	if cfg.Index == "" {
-		cfg.Index = ConfigDefault.Index
-	}
-	if !strings.HasPrefix(cfg.Index, "/") {
-		cfg.Index = "/" + cfg.Index
-	}
-	if cfg.NotFoundFile != "" && !strings.HasPrefix(cfg.NotFoundFile, "/") {
-		cfg.NotFoundFile = "/" + cfg.NotFoundFile
+	// Override config if provided
+	if len(config) > 0 {
+		cfg = config[0]
+
+		// Set default values
+		if cfg.Index == "" {
+			cfg.Index = ConfigDefault.Index
+		}
+		if !strings.HasPrefix(cfg.Index, "/") {
+			cfg.Index = "/" + cfg.Index
+		}
+		if cfg.NotFoundFile != "" && !strings.HasPrefix(cfg.NotFoundFile, "/") {
+			cfg.NotFoundFile = "/" + cfg.NotFoundFile
+		}
 	}
+
 	if cfg.Root == nil {
 		panic("filesystem: Root cannot be nil")
 	}
@@ -72,7 +75,7 @@ func New(config Config) fiber.Handler {
 	var prefix string
 
 	// Return new handler
-	return func(c *fiber.Ctx) error {
+	return func(c *fiber.Ctx) (err error) {
 		// Don't execute middleware if Next returns true
 		if cfg.Next != nil && cfg.Next(c) {
 			return c.Next()
@@ -96,7 +99,12 @@ func New(config Config) fiber.Handler {
 			path = "/" + path
 		}
 
-		file, err := cfg.Root.Open(path)
+		var (
+			file http.File
+			stat os.FileInfo
+		)
+
+		file, err = cfg.Root.Open(path)
 		if err != nil && os.IsNotExist(err) && cfg.NotFoundFile != "" {
 			file, err = cfg.Root.Open(cfg.NotFoundFile)
 		}
@@ -105,12 +113,11 @@ func New(config Config) fiber.Handler {
 			if os.IsNotExist(err) {
 				return c.Status(fiber.StatusNotFound).Next()
 			}
-			return err
+			return
 		}
 
-		stat, err := file.Stat()
-		if err != nil {
-			return err
+		if stat, err = file.Stat(); err != nil {
+			return
 		}
 
 		// Serve index if path is directory
@@ -129,10 +136,7 @@ func New(config Config) fiber.Handler {
 		// Browse directory if no index found and browsing is enabled
 		if stat.IsDir() {
 			if cfg.Browse {
-				if err := dirList(c, file); err != nil {
-					return err
-				}
-				return nil
+				return dirList(c, file)
 			}
 			return fiber.ErrForbidden
 		}
diff --git a/middleware/filesystem/filesystem_test.go b/middleware/filesystem/filesystem_test.go
index bda939a1..6dcbadbf 100644
--- a/middleware/filesystem/filesystem_test.go
+++ b/middleware/filesystem/filesystem_test.go
@@ -28,6 +28,7 @@ func Test_FileSystem(t *testing.T) {
 
 	app.Use("/spatest", New(Config{
 		Root:         http.Dir("../../.github/testdata/fs"),
+		Index:        "index.html",
 		NotFoundFile: "index.html",
 	}))
 
@@ -99,8 +100,7 @@ func Test_FileSystem(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			req, _ := http.NewRequest("GET", tt.url, nil)
-			resp, err := app.Test(req)
+			resp, err := app.Test(httptest.NewRequest("GET", tt.url, nil))
 			utils.AssertEqual(t, nil, err)
 			utils.AssertEqual(t, tt.statusCode, resp.StatusCode)
 
@@ -128,3 +128,38 @@ func Test_FileSystem_Next(t *testing.T) {
 	utils.AssertEqual(t, nil, err)
 	utils.AssertEqual(t, fiber.StatusNotFound, resp.StatusCode)
 }
+
+func Test_FileSystem_NonGetAndHead(t *testing.T) {
+	app := fiber.New()
+
+	app.Use("/test", New(Config{
+		Root: http.Dir("../../.github/testdata/fs"),
+	}))
+
+	resp, err := app.Test(httptest.NewRequest(fiber.MethodPost, "/test", nil))
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, 404, resp.StatusCode)
+}
+
+func Test_FileSystem_Head(t *testing.T) {
+	app := fiber.New()
+
+	app.Use("/test", New(Config{
+		Root: http.Dir("../../.github/testdata/fs"),
+	}))
+
+	req, _ := http.NewRequest(fiber.MethodHead, "/test", nil)
+	resp, err := app.Test(req)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, 200, resp.StatusCode)
+}
+
+func Test_FileSystem_NoRoot(t *testing.T) {
+	defer func() {
+		utils.AssertEqual(t, "filesystem: Root cannot be nil", recover())
+	}()
+
+	app := fiber.New()
+	app.Use(New())
+	_, _ = app.Test(httptest.NewRequest(fiber.MethodGet, "/", nil))
+}