authorize: support arbitrary jwt claims (#2102) (#2106)

* authorize: support arbitrary jwt claims

* remove dead code

Co-authored-by: Caleb Doxsey <cdoxsey@pomerium.com>

authorize/evaluator/opa/policy/authz.rego

@@ -7,6 +7,7 @@ five_minutes := (time.now_ns() / 1e9) + (60 * 5)
 
 # databroker versions to know which version of the data was evaluated
 databroker_server_version := data.databroker_server_version
+
 databroker_record_version := data.databroker_record_version
 
 route_policy_idx := first_allowed_route_policy_idx(input.http.url)
@@ -222,7 +223,7 @@ jwt_payload_groups = v {
 	true
 }
 
-jwt_claims := [
+base_jwt_claims := [
 	["iss", jwt_payload_iss],
 	["aud", jwt_payload_aud],
 	["jti", jwt_payload_jti],
@@ -234,6 +235,25 @@ jwt_claims := [
 	["groups", jwt_payload_groups],
 ]
 
+additional_jwt_claims := [[k, v] |
+	some header_name
+	claim_key := data.jwt_claim_headers[header_name]
+
+	# exclude base_jwt_claims
+	count([1 |
+		[xk, xv] := base_jwt_claims[_]
+		xk == claim_key
+	]) == 0
+
+	# the claim value can come from session claims or user claims
+	claim_value := object.get(session.claims, claim_key, object.get(user.claims, claim_key, null))
+
+	k := claim_key
+	v := get_header_string_value(claim_value)
+]
+
+jwt_claims := array.concat(base_jwt_claims, additional_jwt_claims)
+
 jwt_payload = {key: value |
 	# use a comprehension over an array to remove nil values
 	[key, value] := jwt_claims[_]

authorize/evaluator/opa_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 	"gopkg.in/square/go-jose.v2"
@@ -41,7 +42,7 @@ func TestOPA(t *testing.T) {
 		require.NoError(t, err)
 		store := NewStoreFromProtos(math.MaxUint64, data...)
 		store.UpdateIssuer("authenticate.example.com")
-		store.UpdateJWTClaimHeaders(config.NewJWTClaimHeaders("email", "groups", "user"))
+		store.UpdateJWTClaimHeaders(config.NewJWTClaimHeaders("email", "groups", "user", "CUSTOM_KEY"))
 		store.UpdateRoutePolicies(policies)
 		store.UpdateSigningKey(privateJWK)
 		r := rego.New(
@@ -199,10 +200,29 @@ func TestOPA(t *testing.T) {
 					IdToken: &session.IDToken{
 						IssuedAt: timestamppb.New(time.Date(2021, 2, 1, 1, 1, 1, 1, time.UTC)),
 					},
+					Claims: map[string]*structpb.ListValue{
+						"CUSTOM_KEY": {
+							Values: []*structpb.Value{
+								structpb.NewStringValue("FROM_SESSION"),
+							},
+						},
+						"email": {
+							Values: []*structpb.Value{
+								structpb.NewStringValue("value"),
+							},
+						},
+					},
 				},
 				&user.User{
 					Id:    "user1",
 					Email: "a@example.com",
+					Claims: map[string]*structpb.ListValue{
+						"CUSTOM_KEY": {
+							Values: []*structpb.Value{
+								structpb.NewStringValue("FROM_USER"),
+							},
+						},
+					},
 				},
 				&directory.User{
 					Id:       "user1",
@@ -215,15 +235,16 @@ func TestOPA(t *testing.T) {
 				},
 			)
 			assert.Equal(t, M{
-				"aud":    "from.example.com",
+				"aud":        "from.example.com",
-				"iss":    "authenticate.example.com",
+				"iss":        "authenticate.example.com",
-				"jti":    "session1",
+				"jti":        "session1",
-				"iat":    1612141261.0,
+				"iat":        1612141261.0,
-				"exp":    1609462861.0,
+				"exp":        1609462861.0,
-				"sub":    "user1",
+				"sub":        "user1",
-				"user":   "user1",
+				"user":       "user1",
-				"email":  "a@example.com",
+				"email":      "a@example.com",
-				"groups": A{"group1", "group1name"},
+				"groups":     A{"group1", "group1name"},
+				"CUSTOM_KEY": "FROM_SESSION",
 			}, payload)
 		})
 	})