From 5e7e96af01ecc53547b9315765540fd3a8eaca9b Mon Sep 17 00:00:00 2001
From: mutantmonkey <mutantmonkey@mutantmonkey.in>
Date: Sun, 4 Oct 2015 14:58:00 -0700
Subject: [PATCH] add support for some security headers

This commit adds support for Content-Security-Policy and
X-Frame-Options using the ContentSecurityPolicy middleware.
---
 csp.go       | 40 ++++++++++++++++++++++++++++++++++++++++
 csp_test.go  | 38 ++++++++++++++++++++++++++++++++++++++
 fileserve.go |  2 ++
 server.go    | 34 +++++++++++++++++++++++++---------
 4 files changed, 105 insertions(+), 9 deletions(-)
 create mode 100644 csp.go
 create mode 100644 csp_test.go

diff --git a/csp.go b/csp.go
new file mode 100644
index 0000000..ac68d1a
--- /dev/null
+++ b/csp.go
@@ -0,0 +1,40 @@
+package main
+
+import (
+	"net/http"
+)
+
+const (
+	cspHeader          = "Content-Security-Policy"
+	frameOptionsHeader = "X-Frame-Options"
+)
+
+type csp struct {
+	h    http.Handler
+	opts CSPOptions
+}
+
+type CSPOptions struct {
+	policy string
+	frame  string
+}
+
+func (c csp) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// only add a CSP if one is not already set
+	if existing := w.Header().Get(cspHeader); existing == "" {
+		w.Header().Add(cspHeader, c.opts.policy)
+	}
+
+	w.Header().Set(frameOptionsHeader, c.opts.frame)
+
+	c.h.ServeHTTP(w, r)
+}
+
+func ContentSecurityPolicy(o CSPOptions) func(http.Handler) http.Handler {
+	fn := func(h http.Handler) http.Handler {
+		return csp{h, o}
+	}
+	return fn
+}
+
+// vim:set ts=8 sw=8 noet:
diff --git a/csp_test.go b/csp_test.go
new file mode 100644
index 0000000..ae4c6db
--- /dev/null
+++ b/csp_test.go
@@ -0,0 +1,38 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/zenazn/goji"
+)
+
+var testCSPHeaders = map[string]string{
+	"Content-Security-Policy": "default-src 'none'; style-src 'self';",
+	"X-Frame-Options":         "SAMEORIGIN",
+}
+
+func TestContentSecurityPolicy(t *testing.T) {
+	w := httptest.NewRecorder()
+
+	req, err := http.NewRequest("GET", "/", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	goji.Use(ContentSecurityPolicy(CSPOptions{
+		policy: testCSPHeaders["Content-Security-Policy"],
+		frame:  testCSPHeaders["X-Frame-Options"],
+	}))
+
+	goji.DefaultMux.ServeHTTP(w, req)
+
+	for k, v := range testCSPHeaders {
+		if w.HeaderMap[k][0] != v {
+			t.Fatalf("%s header did not match expected value set by middleware", k)
+		}
+	}
+}
+
+// vim:set ts=8 sw=8 noet:
diff --git a/fileserve.go b/fileserve.go
index e1d2e16..e3fd5f0 100644
--- a/fileserve.go
+++ b/fileserve.go
@@ -26,6 +26,8 @@ func fileServeHandler(c web.C, w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	w.Header().Set("Content-Security-Policy", Config.fileContentSecurityPolicy)
+
 	http.ServeFile(w, r, filePath)
 }
 
diff --git a/server.go b/server.go
index 2f314ec..e87a3ec 100644
--- a/server.go
+++ b/server.go
@@ -19,15 +19,18 @@ import (
 )
 
 var Config struct {
-	bind          string
-	filesDir      string
-	metaDir       string
-	noLogs        bool
-	allowHotlink  bool
-	siteName      string
-	siteURL       string
-	fastcgi       bool
-	remoteUploads bool
+	bind                      string
+	filesDir                  string
+	metaDir                   string
+	noLogs                    bool
+	allowHotlink              bool
+	siteName                  string
+	siteURL                   string
+	fastcgi                   bool
+	remoteUploads             bool
+	contentSecurityPolicy     string
+	fileContentSecurityPolicy string
+	xFrameOptions             string
 }
 
 var Templates = make(map[string]*pongo2.Template)
@@ -37,6 +40,11 @@ var timeStarted time.Time
 var timeStartedStr string
 
 func setup() {
+	goji.Use(ContentSecurityPolicy(CSPOptions{
+		policy: Config.contentSecurityPolicy,
+		frame:  Config.xFrameOptions,
+	}))
+
 	if Config.noLogs {
 		goji.Abandon(middleware.Logger)
 	}
@@ -126,6 +134,14 @@ func main() {
 		"serve through fastcgi")
 	flag.BoolVar(&Config.remoteUploads, "remoteuploads", false,
 		"enable remote uploads")
+	flag.StringVar(&Config.contentSecurityPolicy, "contentSecurityPolicy",
+		"default-src 'self'; img-src 'self' data:; referrer none;",
+		"value of default Content-Security-Policy header")
+	flag.StringVar(&Config.fileContentSecurityPolicy, "fileContentSecurityPolicy",
+		"default-src 'none'; img-src 'self'; object-src 'self'; media-src 'self'; sandbox; referrer none;",
+		"value of Content-Security-Policy header for file access")
+	flag.StringVar(&Config.xFrameOptions, "xFrameOptions", "SAMEORIGIN",
+		"value of X-Frame-Options header")
 	flag.Parse()
 
 	setup()