From 39bb999db65fdac15d41c321cc66467a9b6c93c4 Mon Sep 17 00:00:00 2001
From: mutantmonkey <mutantmonkey@mutantmonkey.in>
Date: Fri, 3 Jun 2016 22:49:01 -0700
Subject: [PATCH] Add ability to set arbitrary headers

This is useful if you want to add headers for things like HTTP Strict
Transport Security or HTTP Public Key Pinning.
---
 headers.go     | 27 +++++++++++++++++++++++++++
 server.go      | 16 ++++++++++++++++
 server_test.go | 18 ++++++++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 headers.go

diff --git a/headers.go b/headers.go
new file mode 100644
index 0000000..918a6dd
--- /dev/null
+++ b/headers.go
@@ -0,0 +1,27 @@
+package main
+
+import (
+	"net/http"
+	"strings"
+)
+
+type addheaders struct {
+	h       http.Handler
+	headers []string
+}
+
+func (a addheaders) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	for _, header := range a.headers {
+		headerSplit := strings.SplitN(header, ": ", 2)
+		w.Header().Add(headerSplit[0], headerSplit[1])
+	}
+
+	a.h.ServeHTTP(w, r)
+}
+
+func AddHeaders(headers []string) func(http.Handler) http.Handler {
+	fn := func(h http.Handler) http.Handler {
+		return addheaders{h, headers}
+	}
+	return fn
+}
diff --git a/server.go b/server.go
index 535da6d..8c129b6 100644
--- a/server.go
+++ b/server.go
@@ -10,6 +10,7 @@ import (
 	"os"
 	"regexp"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/GeertJohan/go.rice"
@@ -20,6 +21,17 @@ import (
 	"github.com/zenazn/goji/web/middleware"
 )
 
+type headerList []string
+
+func (h *headerList) String() string {
+	return strings.Join(*h, ",")
+}
+
+func (h *headerList) Set(value string) error {
+	*h = append(*h, value)
+	return nil
+}
+
 var Config struct {
 	bind                      string
 	filesDir                  string
@@ -40,6 +52,7 @@ var Config struct {
 	remoteUploads             bool
 	authFile                  string
 	remoteAuthFile            string
+	addHeaders                headerList
 }
 
 var Templates = make(map[string]*pongo2.Template)
@@ -69,6 +82,7 @@ func setup() *web.Mux {
 		policy: Config.contentSecurityPolicy,
 		frame:  Config.xFrameOptions,
 	}))
+	mux.Use(AddHeaders(Config.addHeaders))
 
 	if Config.authFile != "" {
 		mux.Use(UploadAuth(AuthOptions{
@@ -205,6 +219,8 @@ func main() {
 		"value of Content-Security-Policy header for file access")
 	flag.StringVar(&Config.xFrameOptions, "xframeoptions", "SAMEORIGIN",
 		"value of X-Frame-Options header")
+	flag.Var(&Config.addHeaders, "addheader",
+		"Add an arbitrary header to the response. This option can be used multiple times.")
 
 	iniflags.Parse()
 
diff --git a/server_test.go b/server_test.go
index 374f229..5cb1d52 100644
--- a/server_test.go
+++ b/server_test.go
@@ -52,6 +52,24 @@ func TestIndex(t *testing.T) {
 	}
 }
 
+func TestAddHeader(t *testing.T) {
+	Config.addHeaders = []string{"Linx-Test: It works!"}
+
+	mux := setup()
+	w := httptest.NewRecorder()
+
+	req, err := http.NewRequest("GET", "/", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mux.ServeHTTP(w, req)
+
+	if w.Header().Get("Linx-Test") != "It works!" {
+		t.Fatal("Header 'Linx-Test: It works!' not found in index response")
+	}
+}
+
 func TestAuthKeys(t *testing.T) {
 	Config.authFile = "/dev/null"