From cdf2af6c2c3b10b3c066ae28ff43e7acd5a6e0b4 Mon Sep 17 00:00:00 2001
From: David Sharnoff <dsharnoff@singlestore.com>
Date: Sun, 5 Feb 2023 23:27:57 -0800
Subject: [PATCH] feat: add CanRefreshTokenInfo to support non-JWT refresh
 tokens  (#244)

* Add an additional, optional, op.Storage interface so that refresh tokens
that are not JWTs do not cause failures when they randomly, sometimes, decrypt
without error

```go
// CanRefreshTokenInfo is an optional additional interface that Storage can support.
// Supporting CanRefreshTokenInfo is required to be able to revoke a refresh token that
// does not happen to also be a JWTs work properly.
type CanRefreshTokenInfo interface {
        // GetRefreshTokenInfo must return oidc.ErrInvalidRefreshToken when presented
	// with a token that is not a refresh token.
	GetRefreshTokenInfo(ctx context.Context, clientID string, token string) (userID string, tokenID string, err error)
}
```

* add comment suggested in code review

* review feedback: return an error defined in op rather than adding a new error to oidc

* move ErrInvalidRefresToken to op/storage.go
---
 NEXT_RELEASE.md            |  1 +
 pkg/op/storage.go          | 12 ++++++++++++
 pkg/op/token_revocation.go | 28 ++++++++++++++++++++++++----
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/NEXT_RELEASE.md b/NEXT_RELEASE.md
index d4f9e71..91f7f5d 100644
--- a/NEXT_RELEASE.md
+++ b/NEXT_RELEASE.md
@@ -3,4 +3,5 @@
 
 - Add `rp/RelyingParty.GetRevokeEndpoint`
 - Rename `op/OpStorage.GetKeyByIDAndUserID` to `op/OpStorage.GetKeyByIDAndClientID` 
+- Add `CanRefreshTokenInfo` (`GetRefreshTokenInfo()`) to `op.Storage`
 
diff --git a/pkg/op/storage.go b/pkg/op/storage.go
index 2b3c93f..153cd21 100644
--- a/pkg/op/storage.go
+++ b/pkg/op/storage.go
@@ -2,6 +2,7 @@ package op
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"gopkg.in/square/go-jose.v2"
@@ -50,6 +51,17 @@ type AuthStorage interface {
 	GetKeySet(context.Context) (*jose.JSONWebKeySet, error)
 }
 
+// CanRefreshTokenInfo is an optional additional interface that Storage can support.
+// Supporting CanRefreshTokenInfo is required to be able to (revoke) a refresh token that
+// is neither an encrypted string of <tokenID>:<userID> nor a JWT.
+type CanRefreshTokenInfo interface {
+	// GetRefreshTokenInfo must return ErrInvalidRefreshToken when presented
+	// with a token that is not a refresh token.
+	GetRefreshTokenInfo(ctx context.Context, clientID string, token string) (userID string, tokenID string, err error)
+}
+
+var ErrInvalidRefreshToken = errors.New("invalid_refresh_token")
+
 type ClientCredentialsStorage interface {
 	ClientCredentialsTokenRequest(ctx context.Context, clientID string, scopes []string) (TokenRequest, error)
 }
diff --git a/pkg/op/token_revocation.go b/pkg/op/token_revocation.go
index db5eea8..9dd0295 100644
--- a/pkg/op/token_revocation.go
+++ b/pkg/op/token_revocation.go
@@ -2,6 +2,7 @@ package op
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"net/url"
 	"strings"
@@ -31,14 +32,33 @@ func revocationHandler(revoker Revoker) func(http.ResponseWriter, *http.Request)
 }
 
 func Revoke(w http.ResponseWriter, r *http.Request, revoker Revoker) {
-	token, _, clientID, err := ParseTokenRevocationRequest(r, revoker)
+	token, tokenTypeHint, clientID, err := ParseTokenRevocationRequest(r, revoker)
 	if err != nil {
 		RevocationRequestError(w, r, err)
 		return
 	}
-	tokenID, subject, ok := getTokenIDAndSubjectForRevocation(r.Context(), revoker, token)
-	if ok {
-		token = tokenID
+	var subject string
+	doDecrypt := true
+	if canRefreshInfo, ok := revoker.Storage().(CanRefreshTokenInfo); ok && tokenTypeHint != "access_token" {
+		userID, tokenID, err := canRefreshInfo.GetRefreshTokenInfo(r.Context(), clientID, token)
+		if err != nil {
+			// An invalid refresh token means that we'll try other things (leaving doDecrypt==true)
+			if !errors.Is(err, ErrInvalidRefreshToken) {
+				RevocationRequestError(w, r, oidc.ErrServerError().WithParent(err))
+				return
+			}
+		} else {
+			token = tokenID
+			subject = userID
+			doDecrypt = false
+		}
+	}
+	if doDecrypt {
+		tokenID, userID, ok := getTokenIDAndSubjectForRevocation(r.Context(), revoker, token)
+		if ok {
+			token = tokenID
+			subject = userID
+		}
 	}
 	if err := revoker.Storage().RevokeToken(r.Context(), token, subject, clientID); err != nil {
 		RevocationRequestError(w, r, err)