[KEYCLOAK-11077] fix refresh token stores

doc.go

@@ -391,7 +391,7 @@ type RequestScope struct {
 // the default practice of a encrypted cookie
 type storage interface {
 	// Set the token to the store
-	Set(string, string) error
+	Set(string, string, time.Duration) error
 	// Get retrieves a token from the store
 	Get(string) (string, error)
 	// Delete removes a key from the store

handlers.go

@@ -32,6 +32,7 @@ import (
 	"time"
 
 	"github.com/coreos/go-oidc/oauth2"
+	"github.com/coreos/go-oidc/oidc"
 
 	"github.com/go-chi/chi"
 	"go.uber.org/zap"
@@ -196,19 +197,24 @@ func (r *oauthProxy) oauthCallbackHandler(w http.ResponseWriter, req *http.Reque
 		// drop in the access token - cookie expiration = access token
 		r.dropAccessTokenCookie(req, w, accessToken, r.getAccessCookieExpiration(token, resp.RefreshToken))
 
+		var expiration time.Duration
+		// notes: not all idp refresh tokens are readable, google for example, so we attempt to decode into
+		// a jwt and if possible extract the expiration, else we default to 10 days
+		var ident *oidc.Identity
+
+		if _, ident, err = parseToken(resp.RefreshToken); err != nil {
+			expiration = 0
+		} else {
+			expiration = time.Until(ident.ExpiresAt)
+		}
+
 		switch r.useStore() {
 		case true:
-			if err = r.StoreRefreshToken(token, encrypted); err != nil {
+			if err = r.StoreRefreshToken(token, encrypted, expiration); err != nil {
 				r.log.Warn("failed to save the refresh token in the store", zap.Error(err))
 			}
 		default:
-			// notes: not all idp refresh tokens are readable, google for example, so we attempt to decode into
-			// a jwt and if possible extract the expiration, else we default to 10 days
-			if _, ident, err := parseToken(resp.RefreshToken); err != nil {
-				r.dropRefreshTokenCookie(req, w, encrypted, 0)
-			} else {
-				r.dropRefreshTokenCookie(req, w, encrypted, time.Until(ident.ExpiresAt))
-			}
+			r.dropRefreshTokenCookie(req, w, encrypted, expiration)
 		}
 	} else {
 		r.dropAccessTokenCookie(req, w, accessToken, time.Until(identity.ExpiresAt))
@@ -300,7 +306,7 @@ func (r *oauthProxy) logoutHandler(w http.ResponseWriter, req *http.Request) {
 		if k == "redirect" {
 			redirectURL = req.URL.Query().Get("redirect")
 			if redirectURL == "" {
-				// than we can default to redirection url
+				// then we can default to redirection url
 				redirectURL = strings.TrimSuffix(r.config.RedirectionURL, "/oauth/callback")
 			}
 		}

middleware.go

@@ -156,7 +156,7 @@ func (r *oauthProxy) authenticationMiddleware() func(http.Handler) http.Handler
 						zap.String("email", user.email))
 
 					// step: check if the user has refresh token
-					refresh, encrypted, err := r.retrieveRefreshToken(req.WithContext(ctx), user)
+					refresh, _, err := r.retrieveRefreshToken(req.WithContext(ctx), user)
 					if err != nil {
 						r.log.Error("unable to find a refresh token for user",
 							zap.String("client_ip", clientIP),
@@ -231,20 +231,22 @@ func (r *oauthProxy) authenticationMiddleware() func(http.Handler) http.Handler
 							w.WriteHeader(http.StatusInternalServerError)
 							return
 						}
-						r.dropRefreshTokenCookie(req.WithContext(ctx), w, encryptedRefreshToken, refreshExpiresIn)
-					}
 
-					if r.useStore() {
-						go func(old, new jose.JWT, encrypted string) {
-							if err := r.DeleteRefreshToken(old); err != nil {
-								r.log.Error("failed to remove old token", zap.Error(err))
-							}
-							if err := r.StoreRefreshToken(new, encrypted); err != nil {
-								r.log.Error("failed to store refresh token", zap.Error(err))
-								return
-							}
-						}(user.token, token, encrypted)
+						if r.useStore() {
+							go func(old, new jose.JWT, encrypted string) {
+								if err := r.DeleteRefreshToken(old); err != nil {
+									r.log.Error("failed to remove old token", zap.Error(err))
+								}
+								if err := r.StoreRefreshToken(new, encrypted, refreshExpiresIn); err != nil {
+									r.log.Error("failed to store refresh token", zap.Error(err))
+									return
+								}
+							}(user.token, token, encryptedRefreshToken)
+						} else {
+							r.dropRefreshTokenCookie(req.WithContext(ctx), w, encryptedRefreshToken, refreshExpiresIn)
+						}
 					}
+
 					// update the with the new access token and inject into the context
 					user.token = token
 					ctx = context.WithValue(req.Context(), contextScopeName, scope)

server.go

@@ -613,6 +613,7 @@ func (r *oauthProxy) createUpstreamProxy(upstream *url.URL) error {
 	// and for refreshed cookies (htts://github.com/keycloak/keycloak-gatekeeper/pulls/456])
 	proxy.KeepDestinationHeaders = true
 	proxy.Logger = httplog.New(ioutil.Discard, "", 0)
+	proxy.KeepDestinationHeaders = true
 	r.upstream = proxy
 
 	// update the tls configuration of the reverse proxy

store_boltdb.go

@@ -60,7 +60,7 @@ func newBoltDBStore(location *url.URL) (storage, error) {
 }
 
 // Set adds a token to the store
-func (r *boltdbStore) Set(key, value string) error {
+func (r *boltdbStore) Set(key, value string, expiration time.Duration) error {
 	return r.client.Update(func(tx *bbolt.Tx) error {
 		bucket := tx.Bucket([]byte(dbName))
 		if bucket == nil {

store_boltdb_test.go

@@ -65,18 +65,21 @@ func TestNewBoltDBStore(t *testing.T) {
 func TestBoltSet(t *testing.T) {
 	s := newTestBoldDB(t)
 	defer s.close()
-	err := s.store.Set("test", "value")
+	err := s.store.Set("test", "value", 0)
 	assert.NoError(t, err)
 }
 
 func TestBoltGet(t *testing.T) {
 	s := newTestBoldDB(t)
 	defer s.close()
+
 	v, err := s.store.Get("test")
 	assert.NoError(t, err)
 	assert.Empty(t, v)
-	err = s.store.Set("test", "value")
+
+	err = s.store.Set("test", "value", 0)
 	assert.NoError(t, err)
+
 	v, err = s.store.Get("test")
 	assert.NoError(t, err)
 	assert.Equal(t, "value", v)
@@ -87,7 +90,7 @@ func TestBoltDelete(t *testing.T) {
 	value := "value"
 	s := newTestBoldDB(t)
 	defer s.close()
-	err := s.store.Set(keyname, value)
+	err := s.store.Set(keyname, value, 0)
 	assert.NoError(t, err)
 	v, err := s.store.Get(keyname)
 	assert.NoError(t, err)

store_redis.go

@@ -47,8 +47,8 @@ func newRedisStore(location *url.URL) (storage, error) {
 }
 
 // Set adds a token to the store
-func (r redisStore) Set(key, value string) error {
-	if err := r.client.Set(key, value, time.Duration(0)); err.Err() != nil {
+func (r redisStore) Set(key, value string, expiration time.Duration) error {
+	if err := r.client.Set(key, value, expiration); err.Err() != nil {
 		return err.Err()
 	}
 
@@ -62,7 +62,7 @@ func (r redisStore) Get(key string) (string, error) {
 		return "", result.Err()
 	}
 
-	return result.String(), nil
+	return result.Val(), nil
 }
 
 // Delete remove the key

stores.go

@@ -18,6 +18,7 @@ package main
 import (
 	"fmt"
 	"net/url"
+	"time"
 
 	"github.com/coreos/go-oidc/jose"
 	"go.uber.org/zap"
@@ -50,8 +51,8 @@ func (r *oauthProxy) useStore() bool {
 }
 
 // StoreRefreshToken the token to the store
-func (r *oauthProxy) StoreRefreshToken(token jose.JWT, value string) error {
-	return r.store.Set(getHashKey(&token), value)
+func (r *oauthProxy) StoreRefreshToken(token jose.JWT, value string, expiration time.Duration) error {
+	return r.store.Set(getHashKey(&token), value, expiration)
 }
 
 // Get retrieves a token from the store, the key we are using here is the access token