From 12da27ddab9250df322d3e166d0e91e6c26f0d5b Mon Sep 17 00:00:00 2001
From: Daniel Carrillo <daniel.carrillo@gmail.com>
Date: Sat, 2 Apr 2022 18:10:48 +0200
Subject: [PATCH] Add optional secure headers to responses

---
 cmd/whatismyip.go                     | 23 ++++++++++++
 integration-tests/integration_test.go |  1 +
 internal/setting/app.go               | 25 ++++++++-----
 internal/setting/app_test.go          | 54 ++++++++++++++-------------
 4 files changed, 69 insertions(+), 34 deletions(-)

diff --git a/cmd/whatismyip.go b/cmd/whatismyip.go
index 9168747..1642c9c 100644
--- a/cmd/whatismyip.go
+++ b/cmd/whatismyip.go
@@ -17,6 +17,7 @@ import (
 	"github.com/dcarrillo/whatismyip/router"
 
 	"github.com/gin-gonic/gin"
+	"github.com/unrolled/secure"
 )
 
 var (
@@ -138,6 +139,28 @@ func setupEngine() {
 	engine = gin.New()
 	engine.Use(gin.LoggerWithFormatter(httputils.GetLogFormatter))
 	engine.Use(gin.Recovery())
+	if setting.App.EnableSecureHeaders {
+		engine.Use(addSecureHeaders())
+	}
 	_ = engine.SetTrustedProxies(nil)
 	engine.TrustedPlatform = setting.App.TrustedHeader
 }
+
+func addSecureHeaders() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		err := secure.New(secure.Options{
+			BrowserXssFilter:   true,
+			ContentTypeNosniff: true,
+			FrameDeny:          true,
+		}).Process(c.Writer, c.Request)
+		if err != nil {
+			c.Abort()
+			return
+		}
+
+		// Avoid header rewrite if response is a redirection.
+		if status := c.Writer.Status(); status > 300 && status < 399 {
+			c.Abort()
+		}
+	}
+}
diff --git a/integration-tests/integration_test.go b/integration-tests/integration_test.go
index f9edca5..024e988 100644
--- a/integration-tests/integration_test.go
+++ b/integration-tests/integration_test.go
@@ -34,6 +34,7 @@ func buildContainer() testcontainers.ContainerRequest {
 			"-tls-crt", "/tmp/server.pem",
 			"-tls-key", "/tmp/server.key",
 			"-trusted-header", "X-Real-IP",
+			"-enable-secure-headers",
 		},
 		ExposedPorts: []string{"8000:8000", "8001:8001"},
 		WaitingFor:   wait.ForLog("Starting TLS server listening on :8001"),
diff --git a/internal/setting/app.go b/internal/setting/app.go
index 0386bf4..30869b0 100644
--- a/internal/setting/app.go
+++ b/internal/setting/app.go
@@ -20,15 +20,16 @@ type serverSettings struct {
 	WriteTimeout time.Duration
 }
 type settings struct {
-	GeodbPath     geodbPath
-	TemplatePath  string
-	BindAddress   string
-	TLSAddress    string
-	TLSCrtPath    string
-	TLSKeyPath    string
-	TrustedHeader string
-	Server        serverSettings
-	version       bool
+	GeodbPath           geodbPath
+	TemplatePath        string
+	BindAddress         string
+	TLSAddress          string
+	TLSCrtPath          string
+	TLSKeyPath          string
+	TrustedHeader       string
+	EnableSecureHeaders bool
+	Server              serverSettings
+	version             bool
 }
 
 const defaultAddress = ":8080"
@@ -74,6 +75,12 @@ func Setup(args []string) (output string, err error) {
 		"Trusted request header for remote IP (e.g. X-Real-IP)",
 	)
 	flags.BoolVar(&App.version, "version", false, "Output version information and exit")
+	flags.BoolVar(
+		&App.EnableSecureHeaders,
+		"enable-secure-headers",
+		false,
+		"Add sane security-related headers to every response",
+	)
 
 	err = flags.Parse(args)
 	if err != nil {
diff --git a/internal/setting/app_test.go b/internal/setting/app_test.go
index 7530840..d38c143 100644
--- a/internal/setting/app_test.go
+++ b/internal/setting/app_test.go
@@ -70,12 +70,13 @@ func TestParseFlags(t *testing.T) {
 					City: "/city-path",
 					ASN:  "/asn-path",
 				},
-				TemplatePath:  "",
-				BindAddress:   ":8080",
-				TLSAddress:    "",
-				TLSCrtPath:    "",
-				TLSKeyPath:    "",
-				TrustedHeader: "",
+				TemplatePath:        "",
+				BindAddress:         ":8080",
+				TLSAddress:          "",
+				TLSCrtPath:          "",
+				TLSKeyPath:          "",
+				TrustedHeader:       "",
+				EnableSecureHeaders: false,
 				Server: serverSettings{
 					ReadTimeout:  10 * time.Second,
 					WriteTimeout: 10 * time.Second,
@@ -89,12 +90,13 @@ func TestParseFlags(t *testing.T) {
 					City: "/city-path",
 					ASN:  "/asn-path",
 				},
-				TemplatePath:  "",
-				BindAddress:   ":8001",
-				TLSAddress:    "",
-				TLSCrtPath:    "",
-				TLSKeyPath:    "",
-				TrustedHeader: "",
+				TemplatePath:        "",
+				BindAddress:         ":8001",
+				TLSAddress:          "",
+				TLSCrtPath:          "",
+				TLSKeyPath:          "",
+				TrustedHeader:       "",
+				EnableSecureHeaders: false,
 				Server: serverSettings{
 					ReadTimeout:  10 * time.Second,
 					WriteTimeout: 10 * time.Second,
@@ -111,12 +113,13 @@ func TestParseFlags(t *testing.T) {
 					City: "/city-path",
 					ASN:  "/asn-path",
 				},
-				TemplatePath:  "",
-				BindAddress:   ":8080",
-				TLSAddress:    ":9000",
-				TLSCrtPath:    "/crt-path",
-				TLSKeyPath:    "/key-path",
-				TrustedHeader: "",
+				TemplatePath:        "",
+				BindAddress:         ":8080",
+				TLSAddress:          ":9000",
+				TLSCrtPath:          "/crt-path",
+				TLSKeyPath:          "/key-path",
+				TrustedHeader:       "",
+				EnableSecureHeaders: false,
 				Server: serverSettings{
 					ReadTimeout:  10 * time.Second,
 					WriteTimeout: 10 * time.Second,
@@ -126,19 +129,20 @@ func TestParseFlags(t *testing.T) {
 		{
 			[]string{
 				"-geoip2-city", "/city-path", "-geoip2-asn", "/asn-path",
-				"-trusted-header", "header",
+				"-trusted-header", "header", "-enable-secure-headers",
 			},
 			settings{
 				GeodbPath: geodbPath{
 					City: "/city-path",
 					ASN:  "/asn-path",
 				},
-				TemplatePath:  "",
-				BindAddress:   ":8080",
-				TLSAddress:    "",
-				TLSCrtPath:    "",
-				TLSKeyPath:    "",
-				TrustedHeader: "header",
+				TemplatePath:        "",
+				BindAddress:         ":8080",
+				TLSAddress:          "",
+				TLSCrtPath:          "",
+				TLSKeyPath:          "",
+				TrustedHeader:       "header",
+				EnableSecureHeaders: true,
 				Server: serverSettings{
 					ReadTimeout:  10 * time.Second,
 					WriteTimeout: 10 * time.Second,