Fix #188

CHANGELOG.md

@@ -13,6 +13,8 @@
 
 ### FIXED
 
+- Fixed HEADERS environment variable parsing [GH-188]
+
 ## v0.0.5
 
 ### NEW

docs/reference/readme.md

@@ -410,7 +410,16 @@ Certificate Authority is set when behind-the-ingress service communication uses
 - Environmental Variable: `HEADERS`
 - Config File Key: `headers`
 - Type: map of `strings` key value pairs
-- Example: `X-Content-Type-Options:nosniff,X-Frame-Options:SAMEORIGIN`
+- Examples: 
+    - Comma Separated: 
+        `X-Content-Type-Options:nosniff,X-Frame-Options:SAMEORIGIN`
+    - JSON: `'{"X-Test": "X-Value"}'`
+    - YAML:
+        ```yaml
+        headers:
+            X-Test: X-Value
+        ```
+
 - To disable: `disable:true`
 - Default :
 

internal/config/options.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"strings"
 	"time"
 
 	"github.com/mitchellh/hashstructure"
@@ -121,7 +122,8 @@ type Options struct {
 	SigningKey string `mapstructure:"signing_key"`
 
 	// Headers to set on all proxied requests. Add a 'disable' key map to turn off.
-	Headers map[string]string `mapstructure:"headers"`
+	HeadersEnv string
+	Headers    map[string]string
 
 	// RefreshCooldown limits the rate a user can refresh her session
 	RefreshCooldown time.Duration `mapstructure:"refresh_cooldown"`
@@ -205,6 +207,12 @@ func OptionsFromViper(configFile string) (Options, error) {
 		return o, fmt.Errorf("failed to parse Policy: %s", err)
 	}
 
+	// Parse Headers
+	err = o.parseHeaders()
+	if err != nil {
+		return o, fmt.Errorf("failed to parse Headers: %s", err)
+	}
+
 	if o.Debug {
 		log.SetDebugMode()
 	}
@@ -324,6 +332,39 @@ func (o *Options) parseURLs() error {
 	return nil
 }
 
+// parseHeaders handles unmarshalling any custom headers correctly from the environment or
+// viper's parsed keys
+func (o *Options) parseHeaders() error {
+	var headers map[string]string
+	if o.HeadersEnv != "" {
+
+		// Handle JSON by default via viper
+		if headers = viper.GetStringMapString("HeadersEnv"); len(headers) == 0 {
+
+			// Try to parse "Key1:Value1,Key2:Value2" syntax
+			headerSlice := strings.Split(o.HeadersEnv, ",")
+			for n := range headerSlice {
+				headerFields := strings.SplitN(headerSlice[n], ":", 2)
+				if len(headerFields) == 2 {
+					headers[headerFields[0]] = headerFields[1]
+
+				} else {
+					// Something went wrong
+					return fmt.Errorf("Failed to decode headers environment variable from '%s'", o.HeadersEnv)
+				}
+			}
+
+		}
+		o.Headers = headers
+	} else if viper.IsSet("headers") {
+		if err := viper.UnmarshalKey("headers", &headers); err != nil {
+			return err
+		}
+		o.Headers = headers
+	}
+	return nil
+}
+
 // bindEnvs makes sure viper binds to each env var based on the mapstructure tag
 func (o *Options) bindEnvs() {
 	tagName := `mapstructure`
@@ -337,6 +378,7 @@ func (o *Options) bindEnvs() {
 
 	// Statically bind fields
 	viper.BindEnv("PolicyEnv", "POLICY")
+	viper.BindEnv("HeadersEnv", "HEADERS")
 }
 
 // findPwd returns best guess at current working directory

internal/config/options_test.go

@@ -148,8 +148,10 @@ func Test_bindEnvs(t *testing.T) {
 	os.Clearenv()
 	defer os.Unsetenv("POMERIUM_DEBUG")
 	defer os.Unsetenv("POLICY")
+	defer os.Unsetenv("HEADERS")
 	os.Setenv("POMERIUM_DEBUG", "true")
 	os.Setenv("POLICY", "mypolicy")
+	os.Setenv("HEADERS", `{"X-Custom-1":"foo", "X-Custom-2":"bar"}`)
 	o.bindEnvs()
 	err := viper.Unmarshal(o)
 	if err != nil {
@@ -164,6 +166,9 @@ func Test_bindEnvs(t *testing.T) {
 	if o.PolicyEnv != "mypolicy" {
 		t.Errorf("Failed to bind policy env var to PolicyEnv")
 	}
+	if o.HeadersEnv != `{"X-Custom-1":"foo", "X-Custom-2":"bar"}` {
+		t.Errorf("Failed to bind headers env var to HeadersEnv")
+	}
 }
 
 func Test_parseURLs(t *testing.T) {
@@ -204,6 +209,43 @@ func Test_parseURLs(t *testing.T) {
 
 }
 
+func Test_parseHeaders(t *testing.T) {
+	tests := []struct {
+		name         string
+		want         map[string]string
+		envHeaders   string
+		viperHeaders interface{}
+		wantErr      bool
+	}{
+		{"good env", map[string]string{"X-Custom-1": "foo", "X-Custom-2": "bar"}, `{"X-Custom-1":"foo", "X-Custom-2":"bar"}`, map[string]string{"X": "foo"}, false},
+		{"good env not_json", map[string]string{"X-Custom-1": "foo", "X-Custom-2": "bar"}, `X-Custom-1:foo,X-Custom-2:bar`, map[string]string{"X": "foo"}, false},
+		{"good viper", map[string]string{"X-Custom-1": "foo", "X-Custom-2": "bar"}, "", map[string]string{"X-Custom-1": "foo", "X-Custom-2": "bar"}, false},
+		{"bad env", map[string]string{}, "xyyyy", map[string]string{"X": "foo"}, true},
+		{"bad env not_json", map[string]string{"X-Custom-1": "foo", "X-Custom-2": "bar"}, `X-Custom-1:foo,X-Custom-2bar`, map[string]string{"X": "foo"}, true},
+		{"bad viper", map[string]string{}, "", "notaheaderstruct", true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			o := NewOptions()
+			viper.Set("headers", tt.viperHeaders)
+			viper.Set("HeadersEnv", tt.envHeaders)
+			o.HeadersEnv = tt.envHeaders
+
+			err := o.parseHeaders()
+
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Error condition unexpected: err=%s", err)
+			}
+
+			if !tt.wantErr && !cmp.Equal(tt.want, o.Headers) {
+				t.Errorf("Did get expected headers: %s", cmp.Diff(tt.want, o.Headers))
+			}
+
+		})
+	}
+}
+
 func Test_OptionsFromViper(t *testing.T) {
 	testPolicy := policy.Policy{
 		To:   "https://httpbin.org",