From 58150f48c52e01609cdfe1385813237e36a4e1ac Mon Sep 17 00:00:00 2001
From: Michael Cardenas <mcardenas@salesforce.com>
Date: Mon, 16 Sep 2024 14:38:52 -0700
Subject: [PATCH] Add token within expiry grace period check

---
 client.go                        | 8 +++++++-
 tokencache/cache_token_source.go | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index 8dc418a..e8e781f 100644
--- a/client.go
+++ b/client.go
@@ -15,7 +15,8 @@ import (
 
 const (
 	// ScopeOfflineAccess requests a refresh token
-	ScopeOfflineAccess = "offline_access"
+	ScopeOfflineAccess         = "offline_access"
+	TokenExpirationGracePeriod = time.Duration(30 * time.Second)
 )
 
 type KeySource interface {
@@ -159,6 +160,11 @@ func (t *Token) Valid() bool {
 		t.IDToken != ""
 }
 
+func (t *Token) WithinGracePeriod() bool {
+	gracePeriod := t.Claims.Expiry.Time().Add(-TokenExpirationGracePeriod)
+	return gracePeriod.After(time.Now()) && t.Valid()
+}
+
 // Type of the token
 func (t *Token) Type() string {
 	// only thing we support for now
diff --git a/tokencache/cache_token_source.go b/tokencache/cache_token_source.go
index 00f0642..c84054e 100644
--- a/tokencache/cache_token_source.go
+++ b/tokencache/cache_token_source.go
@@ -87,7 +87,7 @@ func (c *cachingTokenSource) Token(ctx context.Context) (*oidc.Token, error) {
 	}
 
 	var newToken *oidc.Token
-	if token != nil && token.Valid() {
+	if token != nil && token.Valid() && !token.WithinGracePeriod() {
 		return token, nil
 	} else if token != nil && token.RefreshToken != "" {
 		// we have an expired token, try and refresh if we can.