From 063107e76404ddc3562427b066ad90b2ee11d653 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B8rn=20Erik=20Pedersen?=
 <bjorn.erik.pedersen@gmail.com>
Date: Sun, 8 Mar 2020 16:33:15 +0100
Subject: [PATCH] Add HTTP header support for the dev server

Fixes #7031
---
 commands/commandeer.go      |  6 +++-
 commands/server.go          |  4 +++
 config/commonConfig.go      | 59 +++++++++++++++++++++++++++++++++++++
 config/commonConfig_test.go | 24 +++++++++++++++
 4 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/commands/commandeer.go b/commands/commandeer.go
index 3054ffb7412..547bf8bf326 100644
--- a/commands/commandeer.go
+++ b/commands/commandeer.go
@@ -18,6 +18,8 @@ import (
 	"errors"
 	"sync"
 
+	hconfig "github.com/gohugoio/hugo/config"
+
 	"golang.org/x/sync/semaphore"
 
 	"io/ioutil"
@@ -58,7 +60,8 @@ type commandeerHugoState struct {
 type commandeer struct {
 	*commandeerHugoState
 
-	logger *loggers.Logger
+	logger       *loggers.Logger
+	serverConfig *config.Server
 
 	// Currently only set when in "fast render mode". But it seems to
 	// be fast enough that we could maybe just add it for all server modes.
@@ -343,6 +346,7 @@ func (c *commandeer) loadConfig(mustHaveConfigFile, running bool) error {
 
 	cfg.Logger = logger
 	c.logger = logger
+	c.serverConfig = hconfig.DecodeServer(cfg.Cfg)
 
 	createMemFs := config.GetBool("renderToMemory")
 
diff --git a/commands/server.go b/commands/server.go
index 72884749277..a22a7a69a97 100644
--- a/commands/server.go
+++ b/commands/server.go
@@ -355,6 +355,10 @@ func (f *fileServer) createEndpoint(i int) (*http.ServeMux, string, string, erro
 				w.Header().Set("Pragma", "no-cache")
 			}
 
+			for _, header := range f.c.serverConfig.Match(r.RequestURI) {
+				w.Header().Set(header.Key, header.Value)
+			}
+
 			if f.c.fastRenderMode && f.c.buildErr == nil {
 				p := r.RequestURI
 				if strings.HasSuffix(p, "/") || strings.HasSuffix(p, "html") || strings.HasSuffix(p, "htm") {
diff --git a/config/commonConfig.go b/config/commonConfig.go
index ab2cfe80b7f..17d5619bb12 100644
--- a/config/commonConfig.go
+++ b/config/commonConfig.go
@@ -14,8 +14,13 @@
 package config
 
 import (
+	"sort"
 	"strings"
+	"sync"
 
+	"github.com/gohugoio/hugo/common/types"
+
+	"github.com/gobwas/glob"
 	"github.com/gohugoio/hugo/common/herrors"
 	"github.com/mitchellh/mapstructure"
 	"github.com/spf13/cast"
@@ -88,3 +93,57 @@ func DecodeSitemap(prototype Sitemap, input map[string]interface{}) Sitemap {
 
 	return prototype
 }
+
+// Config for the dev server.
+type Server struct {
+	Headers []Headers
+
+	compiledInit sync.Once
+	compiled     []glob.Glob
+}
+
+func (s *Server) Match(pattern string) []types.KeyValueStr {
+	s.compiledInit.Do(func() {
+		for _, h := range s.Headers {
+			s.compiled = append(s.compiled, glob.MustCompile(h.For))
+		}
+	})
+
+	if s.compiled == nil {
+		return nil
+	}
+
+	var matches []types.KeyValueStr
+
+	for i, g := range s.compiled {
+		if g.Match(pattern) {
+			h := s.Headers[i]
+			for k, v := range h.Values {
+				matches = append(matches, types.KeyValueStr{Key: k, Value: cast.ToString(v)})
+			}
+		}
+	}
+
+	sort.Slice(matches, func(i, j int) bool {
+		return matches[i].Key < matches[j].Key
+	})
+
+	return matches
+
+}
+
+type Headers struct {
+	For    string
+	Values map[string]interface{}
+}
+
+func DecodeServer(cfg Provider) *Server {
+	m := cfg.GetStringMap("server")
+	s := &Server{}
+	if m == nil {
+		return s
+	}
+
+	_ = mapstructure.WeakDecode(m, s)
+	return s
+}
diff --git a/config/commonConfig_test.go b/config/commonConfig_test.go
index 281d2b0b6ea..41b2721bc46 100644
--- a/config/commonConfig_test.go
+++ b/config/commonConfig_test.go
@@ -18,6 +18,7 @@ import (
 	"testing"
 
 	"github.com/gohugoio/hugo/common/herrors"
+	"github.com/gohugoio/hugo/common/types"
 
 	qt "github.com/frankban/quicktest"
 
@@ -58,3 +59,26 @@ func TestBuild(t *testing.T) {
 	c.Assert(b.UseResourceCache(nil), qt.Equals, false)
 
 }
+
+func TestServer(t *testing.T) {
+	c := qt.New(t)
+
+	cfg, err := FromConfigString(`[[server.headers]]
+for = "/*.jpg"
+
+[server.headers.values]
+X-Frame-Options = "DENY"
+X-XSS-Protection = "1; mode=block"
+X-Content-Type-Options = "nosniff"
+`, "toml")
+
+	c.Assert(err, qt.IsNil)
+
+	s := DecodeServer(cfg)
+
+	c.Assert(s.Match("/foo.jpg"), qt.DeepEquals, []types.KeyValueStr{
+		{Key: "X-Content-Type-Options", Value: "nosniff"},
+		{Key: "X-Frame-Options", Value: "DENY"},
+		{Key: "X-XSS-Protection", Value: "1; mode=block"}})
+
+}