From 6f833a5a7306d5abd4dc1719b6e794eedd0ff967 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Wed, 23 Apr 2014 13:10:18 -0700
Subject: [PATCH] agent: Improving UI file serving

---
 command/agent/http.go             |  1 +
 command/agent/http_test.go        |  4 +++-
 command/agent/ui_endpoint.go      | 30 +++++++++++++++++++++++++++++-
 command/agent/ui_endpoint_test.go |  4 +++-
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/command/agent/http.go b/command/agent/http.go
index 39e463b80..b9fd97d30 100644
--- a/command/agent/http.go
+++ b/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)
 
diff --git a/command/agent/http_test.go b/command/agent/http_test.go
index 128b91d43..8d4a63d00 100644
--- a/command/agent/http_test.go
+++ b/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 {
diff --git a/command/agent/ui_endpoint.go b/command/agent/ui_endpoint.go
index a76fc16e9..790506822 100644
--- a/command/agent/ui_endpoint.go
+++ b/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"
 	}
+
+	// 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)
-	http.ServeFile(resp, req, path)
+
+	// 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)
 }
diff --git a/command/agent/ui_endpoint_test.go b/command/agent/ui_endpoint_test.go
index 1db1a8b1b..48de9209a 100644
--- a/command/agent/ui_endpoint_test.go
+++ b/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)