add JWT timestamp formatting workaround (#4270)

Rego will sometimes serialize integers to JSON with a decimal point and
exponent. I don't completely understand this behavior.

Add a workaround to headers.rego to convert the JWT "iat" and "exp"
timestamps to a string and back to an integer. This appears to cause
Rego to serialize these values as plain integers.

Add a unit test to verify this behavior. Also add a unit test that will
fail if the Rego behavior changes, making this workaround unnecessary.

authorize/evaluator/headers_evaluator_test.go

@@ -1,16 +1,24 @@
 package evaluator
 
 import (
+	"bytes"
 	"context"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
 	"math"
+	"strconv"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/go-jose/go-jose/v3/jwt"
+	"github.com/open-policy-agent/opa/rego"
 	"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"
 
 	"github.com/pomerium/pomerium/authorize/internal/store"
 	"github.com/pomerium/pomerium/config"
@@ -67,6 +75,8 @@ func TestHeadersEvaluator(t *testing.T) {
 		return e.Evaluate(ctx, input)
 	}
 
+	iat := time.Unix(1686870680, 0)
+
 	t.Run("jwt", func(t *testing.T) {
 		output, err := eval(t,
 			[]proto.Message{
@@ -75,7 +85,7 @@ func TestHeadersEvaluator(t *testing.T) {
 					"name": {Values: []*structpb.Value{
 						structpb.NewStringValue("n1"),
 					}},
-				}},
+				}, IssuedAt: timestamppb.New(iat)},
 			},
 			&HeadersRequest{
 				Issuer:     "from.example.com",
@@ -86,7 +96,29 @@ func TestHeadersEvaluator(t *testing.T) {
 			})
 		require.NoError(t, err)
 
-		rawJWT, err := jwt.ParseSigned(output.Headers.Get("X-Pomerium-Jwt-Assertion"))
+		jwtHeader := output.Headers.Get("X-Pomerium-Jwt-Assertion")
+
+		// Make sure the 'iat' and 'exp' claims can be parsed as an integer. We
+		// need to do some explicit decoding in order to be able to verify
+		// this, as by default json.Unmarshal() will make no distinction
+		// between numeric formats.
+		d := json.NewDecoder(bytes.NewReader(decodeJWSPayload(t, jwtHeader)))
+		d.UseNumber()
+		var jwtPayloadDecoded map[string]interface{}
+		err = d.Decode(&jwtPayloadDecoded)
+		require.NoError(t, err)
+
+		// The 'iat' claim is set from the session store.
+		assert.Equal(t, json.Number("1686870680"), jwtPayloadDecoded["iat"],
+			"unexpected 'iat' timestamp format")
+
+		// The 'exp' claim will vary with the current time, but we can still
+		// use Atoi() to verify that it can be parsed as an integer.
+		exp := string(jwtPayloadDecoded["exp"].(json.Number))
+		_, err = strconv.Atoi(exp)
+		assert.NoError(t, err, "unexpected 'exp' timestamp format")
+
+		rawJWT, err := jwt.ParseSigned(jwtHeader)
 		require.NoError(t, err)
 
 		var claims M
@@ -188,3 +220,31 @@ func TestHeadersEvaluator(t *testing.T) {
 		assert.Equal(t, "Bearer ID_TOKEN", output.Headers.Get("Authorization"))
 	})
 }
+
+func decodeJWSPayload(t *testing.T, jws string) []byte {
+	t.Helper()
+
+	// A compact JWS string should consist of three base64-encoded values,
+	// separated by a '.' character. The payload is the middle one of these.
+	// cf. https://www.rfc-editor.org/rfc/rfc7515#section-7.1
+	parts := strings.Split(jws, ".")
+	require.Equal(t, 3, len(parts))
+	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
+	require.NoError(t, err)
+	return payload
+}
+
+// If this test fails with the message "workaround no longer needed", then the
+// upstream serialization issue in Rego has been fixed, and we should be able
+// to remove the to_number / format_int workaround from headers.rego (and
+// delete this test).
+func TestTimestampWorkaroundStillNeeded(t *testing.T) {
+	now := strconv.FormatInt(time.Now().Unix(), 10)
+	r := rego.New(rego.Query(fmt.Sprintf("json.marshal(%s + 0)", now)))
+	rs, err := r.Eval(context.Background())
+	require.NoError(t, err, "rego evaluation error")
+	require.Equal(t, 1, len(rs))
+	e := rs[0].Expressions
+	require.Equal(t, 1, len(e))
+	assert.NotEqual(t, now, e[0].Value, "workaround no longer needed")
+}

authorize/evaluator/opa/policy/headers.rego

@@ -90,6 +90,10 @@ jwt_payload_exp = v {
 	v = five_minutes
 } else = null
 
+jwt_payload_exp_int = v {
+	v = to_number(format_int(jwt_payload_exp, 10))
+} else = null
+
 jwt_payload_iat = v {
 	# sessions store the issued_at on the id_token
 	v = round(session.id_token.issued_at.seconds)
@@ -98,6 +102,10 @@ jwt_payload_iat = v {
 	v = round(session.issued_at.seconds)
 } else = null
 
+jwt_payload_iat_int = v {
+	v = to_number(format_int(jwt_payload_iat, 10))
+} else = null
+
 jwt_payload_sub = v {
 	v = session.user_id
 } else = ""
@@ -133,8 +141,8 @@ base_jwt_claims := [
 	["iss", jwt_payload_iss],
 	["aud", jwt_payload_aud],
 	["jti", jwt_payload_jti],
-	["exp", jwt_payload_exp],
+	["exp", jwt_payload_exp_int],
-	["iat", jwt_payload_iat],
+	["iat", jwt_payload_iat_int],
 	["sub", jwt_payload_sub],
 	["user", jwt_payload_user],
 	["email", jwt_payload_email],