From 67373f1e33ed38b071028ee7c60d839df9f03630 Mon Sep 17 00:00:00 2001
From: Johny Michelson <johnymichelson@gmail.com>
Date: Sun, 25 Nov 2018 16:22:22 +0100
Subject: [PATCH] Route StaticFS() not found and forbidden to router's
 NoRoute(), add unit test

---
 routergroup.go | 12 ++++++++++--
 routes_test.go | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/routergroup.go b/routergroup.go
index 9cb0b989df..5615a50cd6 100644
--- a/routergroup.go
+++ b/routergroup.go
@@ -185,11 +185,19 @@ func (group *RouterGroup) StaticFS(relativePath string, fs http.FileSystem) IRou
 func (group *RouterGroup) createStaticHandler(relativePath string, fs http.FileSystem) HandlerFunc {
 	absolutePath := group.calculateAbsolutePath(relativePath)
 	fileServer := http.StripPrefix(absolutePath, http.FileServer(fs))
-	_, nolisting := fs.(*onlyfilesFS)
+
 	return func(c *Context) {
-		if nolisting {
+		file := c.Param("filepath")
+
+		// Check if file exists and/or if we have permission to access it
+		if _, err := fs.Open(file); err != nil {
 			c.Writer.WriteHeader(http.StatusNotFound)
+			c.handlers = group.engine.allNoRoute
+			// Reset index
+			c.index = -1
+			return
 		}
+
 		fileServer.ServeHTTP(c.Writer, c.Request)
 	}
 }
diff --git a/routes_test.go b/routes_test.go
index 60f1c81b6f..c4d59725a9 100644
--- a/routes_test.go
+++ b/routes_test.go
@@ -411,6 +411,21 @@ func TestRouterNotFound(t *testing.T) {
 	assert.Equal(t, http.StatusNotFound, w.Code)
 }
 
+func TestRouterStaticFSNotFound(t *testing.T) {
+	router := New()
+
+	router.StaticFS("/", http.FileSystem(http.Dir("/thisreallydoesntexist/")))
+	router.NoRoute(func(c *Context) {
+		c.String(404, "non existent")
+	})
+
+	w := performRequest(router, "GET", "/nonexistent")
+	assert.Equal(t, "non existent", w.Body.String())
+
+	w = performRequest(router, "HEAD", "/nonexistent")
+	assert.Equal(t, "non existent", w.Body.String())
+}
+
 func TestRouteRawPath(t *testing.T) {
 	route := New()
 	route.UseRawPath = true