From 34fee029d9ad646209b3f861a5f5916a169811f3 Mon Sep 17 00:00:00 2001
From: David Sharnoff <dsharnoff@singlestore.com>
Date: Thu, 10 Nov 2022 17:26:32 -0800
Subject: [PATCH] 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)
}
```
---
 NEXT_RELEASE.md            |  1 +
 pkg/oidc/error.go          |  6 ++++++
 pkg/op/storage.go          |  9 +++++++++
 pkg/op/token_revocation.go | 28 ++++++++++++++++++++++++----
 4 files changed, 40 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/oidc/error.go b/pkg/oidc/error.go
index 5797a59..3aad030 100644
--- a/pkg/oidc/error.go
+++ b/pkg/oidc/error.go
@@ -18,6 +18,7 @@ const (
 	InteractionRequired  errorType = "interaction_required"
 	LoginRequired        errorType = "login_required"
 	RequestNotSupported  errorType = "request_not_supported"
+	InvalidRefreshToken  errorType = "invalid_refresh_token"
 )
 
 var (
@@ -77,6 +78,11 @@ var (
 			ErrorType: RequestNotSupported,
 		}
 	}
+	ErrInvalidRefreshToken = func() *Error {
+		return &Error{
+			ErrorType: InvalidRefreshToken,
+		}
+	}
 )
 
 type Error struct {
diff --git a/pkg/op/storage.go b/pkg/op/storage.go
index 2b3c93f..4cf0aa3 100644
--- a/pkg/op/storage.go
+++ b/pkg/op/storage.go
@@ -50,6 +50,15 @@ 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
+// does not happen to also be JWTs.
+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)
+}
+
 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..0a8c655 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, oidc.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)