Handle shutdown of PAT routine (#506)

* Remove fields used during refactor * Improve goroutine cancelations, server shutdown * Fix lint * Handle shutdown of PAT routine
gogatekeeper · Sep 13, 2024 · 6a5b023 · 6a5b023
1 parent 7c06fdb
commit 6a5b023
Show file tree

Hide file tree

Showing 4 changed files with 164 additions and 127 deletions.
diff --git a/pkg/apperrors/apperrors.go b/pkg/apperrors/apperrors.go
@@ -36,14 +36,14 @@ var (
 	ErrPKCECookieEmpty                 = errors.New("seems that pkce code verifier cookie value is empty string")
 	ErrQueryParamValueMismatch         = errors.New("query param value is not allowed")
 	ErrMissingAuthCode                 = errors.New("missing auth code")
-
-	ErrSessionExpiredVerifyOff  = errors.New("the session has expired and verification switch off")
-	ErrSessionExpiredRefreshOff = errors.New("session expired and access token refreshing is disabled")
-	ErrRefreshTokenNotFound     = errors.New("unable to find refresh token for user")
-	ErrAccTokenRefreshFailure   = errors.New("failed to refresh the access token")
-	ErrEncryptAccToken          = errors.New("unable to encrypt access token")
-	ErrEncryptRefreshToken      = errors.New("failed to encrypt refresh token")
-	ErrEncryptIDToken           = errors.New("unable to encrypt idToken token")
+	ErrInvalidGrantType                = errors.New("invalid grant type is not supported")
+	ErrSessionExpiredVerifyOff         = errors.New("the session has expired and verification switch off")
+	ErrSessionExpiredRefreshOff        = errors.New("session expired and access token refreshing is disabled")
+	ErrRefreshTokenNotFound            = errors.New("unable to find refresh token for user")
+	ErrAccTokenRefreshFailure          = errors.New("failed to refresh the access token")
+	ErrEncryptAccToken                 = errors.New("unable to encrypt access token")
+	ErrEncryptRefreshToken             = errors.New("failed to encrypt refresh token")
+	ErrEncryptIDToken                  = errors.New("unable to encrypt idToken token")
 
 	ErrDelTokFromStore = errors.New("failed to remove old token")
 	ErrSaveTokToStore  = errors.New("failed to store refresh token")
@@ -58,9 +58,10 @@ var (
 	ErrParseRefreshToken              = errors.New("failed to parse refresh token")
 	ErrParseIDToken                   = errors.New("failed to parse id token")
 	ErrParseAccessToken               = errors.New("failed to parse access token")
-	ErrParseIDTokenClaims             = errors.New("faled to parse id token claims")
-	ErrParseAccessTokenClaims         = errors.New("faled to parse access token claims")
-	ErrParseRefreshTokenClaims        = errors.New("faled to parse refresh token claims")
+	ErrParseIDTokenClaims             = errors.New("failed to parse id token claims")
+	ErrParseAccessTokenClaims         = errors.New("failed to parse access token claims")
+	ErrParseRefreshTokenClaims        = errors.New("failed to parse refresh token claims")
+	ErrPATTokenFetch                  = errors.New("failed to get PAT token")
 
 	ErrAccTokenVerifyFailure   = errors.New("access token failed verification")
 	ErrTokenSignature          = errors.New("invalid token signature")
@@ -78,6 +79,10 @@ var (
 	ErrHmacHeaderEmpty = errors.New("request HMAC header empty")
 	ErrHmacMismatch    = errors.New("received HMAC header and calculated HMAC does not match")
 
+	ErrStartMainHTTP     = errors.New("failed to start main http service")
+	ErrStartRedirectHTTP = errors.New("failed to start http redirect service")
+	ErrStartAdminHTTP    = errors.New("failed to start admin service")
+
 	// config errors
 
 	ErrNoRedirectsWithEnableRefreshTokensInvalid = errors.New("no-redirects true cannot be enabled with refresh tokens")

diff --git a/pkg/keycloak/proxy/misc.go b/pkg/keycloak/proxy/misc.go
@@ -19,11 +19,11 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"os"
 	"strings"
 	"time"
 
 	"github.com/Nerzal/gocloak/v12"
+	"github.com/cenkalti/backoff/v4"
 	oidc3 "github.com/coreos/go-oidc/v3/oidc"
 	"github.com/go-jose/go-jose/v4/jwt"
 	"github.com/gogatekeeper/gatekeeper/pkg/apperrors"
@@ -36,8 +36,67 @@ import (
 	"go.uber.org/zap"
 )
 
-//nolint:cyclop
 func getPAT(
+	ctx context.Context,
+	clientID string,
+	clientSecret string,
+	realm string,
+	openIDProviderTimeout time.Duration,
+	grantType string,
+	idpClient *gocloak.GoCloak,
+	forwardingUsername string,
+	forwardingPassword string,
+) (*gocloak.JWT, *jwt.Claims, error) {
+	cntx, cancel := context.WithTimeout(
+		ctx,
+		openIDProviderTimeout,
+	)
+	defer cancel()
+
+	var token *gocloak.JWT
+	var err error
+
+	switch grantType {
+	case configcore.GrantTypeClientCreds:
+		token, err = idpClient.LoginClient(
+			cntx,
+			clientID,
+			clientSecret,
+			realm,
+		)
+	case configcore.GrantTypeUserCreds:
+		token, err = idpClient.Login(
+			cntx,
+			clientID,
+			clientSecret,
+			realm,
+			forwardingUsername,
+			forwardingPassword,
+		)
+	default:
+		return nil, nil, apperrors.ErrInvalidGrantType
+	}
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	parsedToken, err := jwt.ParseSigned(token.AccessToken, constant.SignatureAlgs[:])
+	if err != nil {
+		return nil, nil, err
+	}
+
+	stdClaims := &jwt.Claims{}
+	err = parsedToken.UnsafeClaimsWithoutVerification(stdClaims)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return token, stdClaims, err
+}
+
+func refreshPAT(
+	ctx context.Context,
 	logger *zap.Logger,
 	pat *PAT,
 	clientID string,
@@ -52,8 +111,7 @@ func getPAT(
 	forwardingUsername string,
 	forwardingPassword string,
 	done chan bool,
-) {
-	retry := 0
+) error {
 	initialized := false
 	grantType := configcore.GrantTypeClientCreds
 
@@ -62,57 +120,44 @@ func getPAT(
 	}
 
 	for {
-		if retry > 0 {
-			logger.Info(
-				"retrying fetching PAT token",
-				zap.Int("retry", retry),
-			)
-		}
-
-		ctx, cancel := context.WithTimeout(
-			context.Background(),
-			openIDProviderTimeout,
-		)
-
 		var token *gocloak.JWT
-		var err error
-
-		switch grantType {
-		case configcore.GrantTypeClientCreds:
-			token, err = idpClient.LoginClient(
-				ctx,
-				clientID,
-				clientSecret,
-				realm,
-			)
-		case configcore.GrantTypeUserCreds:
-			token, err = idpClient.Login(
-				ctx,
+		var claims *jwt.Claims
+		operation := func() error {
+			var err error
+			pCtx, cancel := context.WithCancel(ctx)
+			defer cancel()
+			token, claims, err = getPAT(
+				pCtx,
 				clientID,
 				clientSecret,
 				realm,
+				openIDProviderTimeout,
+				grantType,
+				idpClient,
 				forwardingUsername,
 				forwardingPassword,
 			)
-		default:
+			return err
+		}
+
+		notify := func(err error, delay time.Duration) {
 			logger.Error(
-				"Chosen grant type is not supported",
-				zap.String("grant_type", grantType),
+				err.Error(),
+				zap.Duration("retry after", delay),
 			)
-			os.Exit(1)
 		}
 
-		if err != nil {
-			retry++
-			logger.Error("problem getting PAT token", zap.Error(err))
-
-			if retry >= patRetryCount {
-				cancel()
-				os.Exit(1)
-			}
+		bom := backoff.WithMaxRetries(
+			backoff.NewConstantBackOff(patRetryInterval),
+			uint64(patRetryCount),
+		)
+		boCtx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		box := backoff.WithContext(bom, boCtx)
+		err := backoff.RetryNotify(operation, box, notify)
 
-			<-time.After(patRetryInterval)
-			continue
+		if err != nil {
+			return err
 		}
 
 		pat.m.Lock()
@@ -121,37 +166,25 @@ func getPAT(
 
 		if !initialized {
 			done <- true
+			initialized = true
 		}
 
-		initialized = true
-
-		parsedToken, err := jwt.ParseSigned(token.AccessToken, constant.SignatureAlgs[:])
-		if err != nil {
-			retry++
-			logger.Error("failed to parse the access token", zap.Error(err))
-			<-time.After(patRetryInterval)
-			continue
-		}
-
-		stdClaims := &jwt.Claims{}
-		err = parsedToken.UnsafeClaimsWithoutVerification(stdClaims)
-		if err != nil {
-			retry++
-			logger.Error("unable to parse access token for claims", zap.Error(err))
-			<-time.After(patRetryInterval)
-			continue
-		}
-
-		retry = 0
-		expiration := stdClaims.Expiry.Time()
+		expiration := claims.Expiry.Time()
 		refreshIn := utils.GetWithin(expiration, constant.PATRefreshInPercent)
 
 		logger.Info(
-			"waiting for expiration of access token",
+			"waiting for access token expiration",
 			zap.Float64("refresh_in", refreshIn.Seconds()),
 		)
 
-		<-time.After(refreshIn)
+		refreshTimer := time.NewTimer(refreshIn)
+		select {
+		case <-ctx.Done():
+			logger.Info("shutdown PAT refresh routine")
+			refreshTimer.Stop()
+			return nil
+		case <-refreshTimer.C:
+		}
 	}
 }