From ba90cb401d8669aedf86573d1e305ef5648ae711 Mon Sep 17 00:00:00 2001
From: Yuxuan 'fishy' Wang <yuxuan.wang@reddit.com>
Date: Fri, 21 Feb 2020 20:51:30 -0800
Subject: [PATCH] edgecontext: Use reddit's fork of jwt-go for key rotation

This fixes https://github.com/reddit/baseplate.go/issues/73.
---
 edgecontext/validator.go | 76 ++++++++++++----------------------------
 go.mod                   |  2 ++
 go.sum                   |  4 +--
 3 files changed, 26 insertions(+), 56 deletions(-)

diff --git a/edgecontext/validator.go b/edgecontext/validator.go
index f2cac9772..94a138e43 100644
--- a/edgecontext/validator.go
+++ b/edgecontext/validator.go
@@ -10,42 +10,13 @@ import (
 	"github.com/reddit/baseplate.go/secrets"
 )
 
-type keysType []*rsa.PublicKey
+type keysType = []*rsa.PublicKey
 
 const (
 	authenticationPubKeySecretPath = "secret/authentication/public-key"
 	jwtAlg                         = "RS256"
 )
 
-// When trying versioned secret with jwt, there are some errors that won't be
-// fixed by the next version of the secret, so we can return early instead of
-// trying all the remaining versions.
-//
-// TODO: We can also get rid of this block when upstream added native support
-// for key rotation.
-var shortCircuitErrors = []uint32{
-	jwt.ValidationErrorMalformed,
-	jwt.ValidationErrorAudience,
-	jwt.ValidationErrorExpired,
-	jwt.ValidationErrorIssuedAt,
-	jwt.ValidationErrorIssuer,
-	jwt.ValidationErrorNotValidYet,
-	jwt.ValidationErrorId,
-	jwt.ValidationErrorClaimsInvalid,
-}
-
-func shouldShortCircutError(err error) bool {
-	var ve jwt.ValidationError
-	if errors.As(err, &ve) {
-		for _, bitmask := range shortCircuitErrors {
-			if ve.Errors&bitmask != 0 {
-				return true
-			}
-		}
-	}
-	return false
-}
-
 // ValidateToken parses and validates a jwt token, and return the decoded
 // AuthenticationToken.
 func ValidateToken(token string) (*AuthenticationToken, error) {
@@ -55,33 +26,30 @@ func ValidateToken(token string) (*AuthenticationToken, error) {
 		return nil, errors.New("no public keys loaded")
 	}
 
-	// TODO: Patch upstream to support key rotation natively:
-	// https://github.com/dgrijalva/jwt-go/pull/372
-	var lastErr error
-	for _, key := range keys {
-		token, err := jwt.ParseWithClaims(
-			token,
-			&AuthenticationToken{},
-			func(_ *jwt.Token) (interface{}, error) {
-				return key, nil
-			},
-		)
-		if err != nil {
-			if shouldShortCircutError(err) {
-				return nil, err
-			}
-			// Try next pubkey.
-			lastErr = err
-			continue
-		}
+	tok, err := jwt.ParseWithClaims(
+		token,
+		&AuthenticationToken{},
+		func(_ *jwt.Token) (interface{}, error) {
+			return keys, nil
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
 
-		if claims, ok := token.Claims.(*AuthenticationToken); ok && token.Valid && token.Method.Alg() == jwtAlg {
-			return claims, nil
-		}
+	if !tok.Valid {
+		return nil, jwt.NewValidationError("invalid token", 0)
+	}
+
+	if tok.Method.Alg() != jwtAlg {
+		return nil, jwt.NewValidationError("wrong signing method", 0)
+	}
 
-		lastErr = jwt.NewValidationError("", 0)
+	if claims, ok := tok.Claims.(*AuthenticationToken); ok {
+		return claims, nil
 	}
-	return nil, lastErr
+
+	return nil, jwt.NewValidationError("invalid token type", 0)
 }
 
 func validatorMiddleware(next secrets.SecretHandlerFunc) secrets.SecretHandlerFunc {
diff --git a/go.mod b/go.mod
index 9f37303ec..4fb5e2503 100644
--- a/go.mod
+++ b/go.mod
@@ -18,3 +18,5 @@ require (
 	gopkg.in/dgrijalva/jwt-go.v3 v3.2.0
 	gopkg.in/fsnotify.v1 v1.4.7
 )
+
+replace gopkg.in/dgrijalva/jwt-go.v3 => github.com/reddit/jwt-go v3.2.1-0.20200222044038-a63f2d40479f+incompatible
diff --git a/go.sum b/go.sum
index 72d44cbc5..60a110ba1 100644
--- a/go.sum
+++ b/go.sum
@@ -48,6 +48,8 @@ github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/reddit/jwt-go v3.2.1-0.20200222044038-a63f2d40479f+incompatible h1:d2fV4H2zMs1kC0dw5N9qbsWW45SsRQSta8IlWEwAG4g=
+github.com/reddit/jwt-go v3.2.1-0.20200222044038-a63f2d40479f+incompatible/go.mod h1:DnRZZdtPlHMhfOZTDM2U49R+PsC3qEV0E+y6rr7Od3o=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
@@ -95,8 +97,6 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/dgrijalva/jwt-go.v3 v3.2.0 h1:N46iQqOtHry7Hxzb9PGrP68oovQmj7EhudNoKHvbOvI=
-gopkg.in/dgrijalva/jwt-go.v3 v3.2.0/go.mod h1:hdNXC2Z9yC029rvsQ/on2ZNQ44Z2XToVhpXXbR+J05A=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=