authorize: rework token substitution in headers (#4456)

Currently Pomerium replaces dynamic set_request_headers tokens
sequentially. As a result, if a replacement value itself contained a
supported "$pomerium" token, Pomerium may treat that as another
replacement, resulting in incorrect output.

This is unlikely to be a problem given the current set of dynamic
tokens, but if we continue to add additional tokens, this will likely
become more of a concern.

To forestall any issues, let's perform all replacements in one pass,
using the os.Expand() method. This does require a slight change to the
syntax, as tokens containing a '.' will need to be wrapped in curly
braces, e.g. ${pomerium.id_token}.

A literal dollar sign can be included by using $$ in the input.

authorize/evaluator/headers_evaluator.go

@@ -4,9 +4,12 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"os"
 
 	envoy_config_cluster_v3 "github.com/envoyproxy/go-control-plane/envoy/config/cluster/v3"
+	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/types"
 
 	"github.com/pomerium/pomerium/authorize/evaluator/opa"
 	"github.com/pomerium/pomerium/authorize/internal/store"
@@ -54,6 +57,48 @@ type HeadersResponse struct {
 	Headers http.Header
 }
 
+var variableSubstitutionFunctionRegoOption = rego.Function2(&rego.Function{
+	Name: "pomerium.variable_substitution",
+	Decl: types.NewFunction(
+		types.Args(
+			types.Named("input_string", types.S),
+			types.Named("replacements",
+				types.NewObject(nil, types.NewDynamicProperty(types.S, types.S))),
+		),
+		types.Named("output", types.S),
+	),
+}, func(bctx rego.BuiltinContext, op1 *ast.Term, op2 *ast.Term) (*ast.Term, error) {
+	inputString, ok := op1.Value.(ast.String)
+	if !ok {
+		return nil, fmt.Errorf("invalid input_string type: %T", op1.Value)
+	}
+
+	replacements, ok := op2.Value.(ast.Object)
+	if !ok {
+		return nil, fmt.Errorf("invalid replacements type: %T", op2.Value)
+	}
+
+	var err error
+	output := os.Expand(string(inputString), func(key string) string {
+		if key == "$" {
+			return "$" // allow a dollar sign to be escaped using $$
+		}
+		r := replacements.Get(ast.StringTerm(key))
+		if r == nil {
+			return ""
+		}
+		s, ok := r.Value.(ast.String)
+		if !ok {
+			err = fmt.Errorf("invalid replacement value type for key %q: %T", key, r.Value)
+		}
+		return string(s)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return ast.StringTerm(output), nil
+})
+
 // A HeadersEvaluator evaluates the headers.rego script.
 type HeadersEvaluator struct {
 	q rego.PreparedEvalQuery
@@ -66,6 +111,7 @@ func NewHeadersEvaluator(ctx context.Context, store *store.Store) (*HeadersEvalu
 		rego.Module("pomerium.headers", opa.HeadersRego),
 		rego.Query("result = data.pomerium.headers"),
 		getGoogleCloudServerlessHeadersRegoOption,
+		variableSubstitutionFunctionRegoOption,
 		store.GetDataBrokerRecordOption(),
 	)
 

authorize/evaluator/headers_evaluator_test.go

@@ -193,9 +193,10 @@ func TestHeadersEvaluator(t *testing.T) {
 				Session:    RequestSession{ID: "s1"},
 				SetRequestHeaders: map[string]string{
 					"X-Custom-Header":         "CUSTOM_VALUE",
-					"X-ID-Token":              "$pomerium.id_token",
-					"X-Access-Token":          "$pomerium.access_token",
-					"Client-Cert-Fingerprint": "$pomerium.client_cert_fingerprint",
+					"X-ID-Token":              "${pomerium.id_token}",
+					"X-Access-Token":          "${pomerium.access_token}",
+					"Client-Cert-Fingerprint": "${pomerium.client_cert_fingerprint}",
+					"Foo":                     "escaped $$dollar sign",
 				},
 				ClientCertificate: ClientCertificateInfo{Leaf: testValidCert},
 			})
@@ -206,6 +207,29 @@ func TestHeadersEvaluator(t *testing.T) {
 		assert.Equal(t, "ACCESS_TOKEN", output.Headers.Get("X-Access-Token"))
 		assert.Equal(t, "ebf421e323e31c3900a7985a16e72c59f45f5a2c15283297567e226b3b17d1a1",
 			output.Headers.Get("Client-Cert-Fingerprint"))
+		assert.Equal(t, "escaped $dollar sign", output.Headers.Get("Foo"))
+	})
+
+	t.Run("set_request_headers no repeated substitution", func(t *testing.T) {
+		output, err := eval(t,
+			[]proto.Message{
+				&session.Session{Id: "s1", IdToken: &session.IDToken{
+					Raw: "$pomerium.access_token",
+				}, OauthToken: &session.OAuthToken{
+					AccessToken: "ACCESS_TOKEN",
+				}},
+			},
+			&HeadersRequest{
+				Issuer:     "from.example.com",
+				ToAudience: "to.example.com",
+				Session:    RequestSession{ID: "s1"},
+				SetRequestHeaders: map[string]string{
+					"X-ID-Token": "${pomerium.id_token}",
+				},
+			})
+		require.NoError(t, err)
+
+		assert.Equal(t, "$pomerium.access_token", output.Headers.Get("X-ID-Token"))
 	})
 
 	t.Run("set_request_headers original behavior", func(t *testing.T) {
@@ -222,7 +246,7 @@ func TestHeadersEvaluator(t *testing.T) {
 				ToAudience: "to.example.com",
 				Session:    RequestSession{ID: "s1"},
 				SetRequestHeaders: map[string]string{
-					"Authorization": "Bearer $pomerium.id_token",
+					"Authorization": "Bearer ${pomerium.id_token}",
 				},
 			})
 		require.NoError(t, err)
@@ -236,7 +260,7 @@ func TestHeadersEvaluator(t *testing.T) {
 				Issuer:     "from.example.com",
 				ToAudience: "to.example.com",
 				SetRequestHeaders: map[string]string{
-					"fingerprint": "$pomerium.client_cert_fingerprint",
+					"fingerprint": "${pomerium.client_cert_fingerprint}",
 				},
 			})
 		require.NoError(t, err)

authorize/evaluator/opa/policy/headers.rego

@@ -219,13 +219,15 @@ client_cert_fingerprint = v {
 } else = ""
 
 set_request_headers = h {
+    replacements := {
+        "pomerium.id_token": session_id_token,
+        "pomerium.access_token": session_access_token,
+		"pomerium.client_cert_fingerprint": client_cert_fingerprint,
+    }
 	h := [[header_name, header_value] |
 		some header_name
-		v1 := input.set_request_headers[header_name]
-		v2 := replace(v1, "$pomerium.id_token", session_id_token)
-		v3 := replace(v2, "$pomerium.access_token", session_access_token)
-		v4 := replace(v3, "$pomerium.client_cert_fingerprint", client_cert_fingerprint)
-		header_value := v4
+		v := input.set_request_headers[header_name]
+		header_value := pomerium.variable_substitution(v, replacements)
 	]
 } else = []
 

config/policy.go

@@ -544,7 +544,7 @@ func (p *Policy) Validate() error {
 
 	if p.SetAuthorizationHeader != "" {
 		log.Warn(context.Background()).Msg("config: set_authorization_header is deprecated, " +
-			"use $pomerium.id_token or $pomerium.access_token in set_request_headers instead")
+			"use ${pomerium.id_token} or ${pomerium.access_token} in set_request_headers instead")
 	}
 
 	return nil