agent: Improving UI file serving

command/agent/http.go

@@ -42,6 +42,7 @@ func NewHTTPServer(agent *Agent, uiDir string, enableDebug bool, logOutput io.Wr
 		mux:      mux,
 		listener: list,
 		logger:   log.New(logOutput, "", log.LstdFlags),
+		uiDir:    uiDir,
 	}
 	srv.registerHandlers(enableDebug)
 

command/agent/http_test.go

@@ -19,7 +19,9 @@ func makeHTTPServer(t *testing.T) (string, *HTTPServer) {
 	conf := nextConfig()
 	dir, agent := makeAgent(t, conf)
 	uiDir := filepath.Join(dir, "ui")
-	println(uiDir)
+	if err := os.Mkdir(uiDir, 755); err != nil {
+		t.Fatalf("err: %v", err)
+	}
 	addr, _ := agent.config.ClientListener(agent.config.Ports.HTTP)
 	server, err := NewHTTPServer(agent, uiDir, true, agent.logOutput, addr.String())
 	if err != nil {

command/agent/ui_endpoint.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"net/http"
+	"os"
 	"path/filepath"
 	"strings"
 	"time"
@@ -15,10 +16,37 @@ func (s *HTTPServer) UiIndex(resp http.ResponseWriter, req *http.Request) {
 		s.logger.Printf("[DEBUG] http: Request %v (%v)", req.URL, time.Now().Sub(start))
 	}()
 
+	// Determine the file path
 	file := strings.TrimPrefix(req.URL.Path, "/ui/")
 	if file == "" {
 		file = "index.html"
 	}
-	path := filepath.Join(s.uiDir, file)
-	http.ServeFile(resp, req, path)
+
+	// Reject if it is relative
+	if strings.Contains(file, "..") {
+		s.logger.Printf("[WARN] http: Invalid file %s", file)
+		resp.WriteHeader(404)
+		return
+	}
+	path := filepath.Join(s.uiDir, file)
+
+	// Attempt to open
+	fh, err := os.Open(path)
+	if err != nil {
+		s.logger.Printf("[WARN] http: Failed to serve file %s: %v", path, err)
+		resp.WriteHeader(404)
+		return
+	}
+	defer fh.Close()
+
+	// Get the file info
+	info, err := fh.Stat()
+	if err != nil {
+		s.logger.Printf("[WARN] http: Failed to stat file %s: %v", path, err)
+		resp.WriteHeader(404)
+		return
+	}
+
+	// Serve the file
+	http.ServeContent(resp, req, path, info.ModTime(), fh)
 }

command/agent/ui_endpoint_test.go

@@ -19,7 +19,9 @@ func TestUiIndex(t *testing.T) {
 
 	// Create file
 	path := filepath.Join(srv.uiDir, "my-file")
-	ioutil.WriteFile(path, []byte("test"), 777)
+	if err := ioutil.WriteFile(path, []byte("test"), 777); err != nil {
+		t.Fatalf("err: %v", err)
+	}
 
 	// Register node
 	req, err := http.NewRequest("GET", "/ui/my-file", nil)