allow specifying Cache-Control header for static files

1 files changed, 23 insertions, 7 deletions

diff --git a/main.go b/main.go
index a290ce7..33d1cee 100644
--- a/main.go
+++ b/main.go
@@ -114,11 +114,21 @@ func main() {
 		return luaOk(l, nil)
 	}
 	routeModule["static"] = func (l *Lua) int {
-		var route, dir string
-		err := l.Scan(&route, &dir)
+		var route, dir, cacheControl string
+		err := l.Scan(&route, &dir, &cacheControl)
 		if err != nil {
 			return luaErr(l, err)
 		}
+		if route == "" {
+			return luaErr(l, errors.New("route must not be empty"))
+		}
+		if dir == "" {
+			return luaErr(l, errors.New("dir must not be empty"))
+		}
+		if !strings.HasSuffix(route, "/") {
+			route = route + "/"
+		}
+		dir = strings.TrimSuffix(dir, "/")
 		// check if route is already registered. otherwise
 		// mux.HandleFunc will panic
 		_, ok := routes[route]
@@ -127,12 +137,18 @@ func main() {
 		}
 		routes[route] = true
 		mu.Lock()
-		mux.Handle(
+		mux.HandleFunc(
 			route,
-			http.StripPrefix(
-				strings.TrimPrefix(route, "GET "),
-				http.FileServer(http.Dir(dir)),
-			),
+			func (w http.ResponseWriter, r *http.Request) {
+				path := strings.Trim(r.URL.Path, "/")
+				route := strings.TrimPrefix(route, "GET ")
+				route = strings.Trim(route, "/")
+				filepath := dir + strings.TrimPrefix(path, route)
+				if cacheControl != "" {
+					w.Header().Set("Cache-Control", cacheControl)
+				}
+				http.ServeFile(w, r, filepath)
+			},
 		)
 		mu.Unlock()
 		return luaOk(l, nil)