From 574c71c40d33c8bbbed19b22821b57b3e084b887 Mon Sep 17 00:00:00 2001
From: Bani Singh <47721811+banikharbanda@users.noreply.github.com>
Date: Mon, 3 Jun 2024 09:58:42 -0400
Subject: [PATCH] Fix refresh expiry code and add tests (#577)

---
 go/pkg/credshelper/credshelper.go      | 12 ++++++++++++
 go/pkg/credshelper/credshelper_test.go | 17 ++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/go/pkg/credshelper/credshelper.go b/go/pkg/credshelper/credshelper.go
index 0376e5019..b3e9ecc1b 100644
--- a/go/pkg/credshelper/credshelper.go
+++ b/go/pkg/credshelper/credshelper.go
@@ -166,6 +166,14 @@ func (c *Credentials) RemoveFromDisk() {
 	}
 }
 
+// refreshStatus checks refresh expiry of credentials in case a manual refresh is required.
+func (c *Credentials) refreshStatus() error {
+	if !c.refreshExp.IsZero() && c.refreshExp.Before(nowFn()) {
+		return fmt.Errorf("credentials cannot be refreshed automatically, manual re-authentication required")
+	}
+	return nil
+}
+
 // Token retrieves an oauth2 token from the external tokensource.
 func (ts *externalTokenSource) Token() (*oauth2.Token, error) {
 	if ts == nil {
@@ -215,6 +223,10 @@ func NewExternalCredentials(credshelper string, credshelperArgs []string, credsF
 			return creds, nil
 		}
 		log.Warningf("Failed to use cached credentials: %v", err)
+		if err = creds.refreshStatus(); err != nil {
+			creds.RemoveFromDisk()
+			return nil, err
+		}
 	}
 	credsOut, err := runCredsHelperCmd(credsHelperCmd)
 	if err != nil {
diff --git a/go/pkg/credshelper/credshelper_test.go b/go/pkg/credshelper/credshelper_test.go
index 9fc124517..3d5e866c8 100644
--- a/go/pkg/credshelper/credshelper_test.go
+++ b/go/pkg/credshelper/credshelper_test.go
@@ -26,7 +26,7 @@ func TestCredentialsHelperCache(t *testing.T) {
 	if err != nil {
 		t.Errorf("failed to create dir for credentials file %q: %v", cf, err)
 	}
-	credsHelperCmd := newReusableCmd("echo", []string{`{"headers":{"hdr":"val"},"token":"testToken", "expiry":""}`})
+	credsHelperCmd := newReusableCmd("echo", []string{`{"headers":{"hdr":"val"},"token":"testToken", "expiry":"","refresh_expiry":""}`})
 	ts := &grpcOauth.TokenSource{
 		TokenSource: oauth2.ReuseTokenSourceWithExpiry(
 			&oauth2.Token{},
@@ -318,6 +318,21 @@ func TestGetRequestMetadata(t *testing.T) {
 	}
 }
 
+func TestRefreshStatus(t *testing.T) {
+	c := Credentials{refreshExp: time.Time{}}
+	if err := c.refreshStatus(); err != nil {
+		t.Errorf("RefreshStatus returned an error when refreshExpiry is zero")
+	}
+	c.refreshExp = time.Now().Add(time.Hour)
+	if err := c.refreshStatus(); err != nil {
+		t.Errorf("RefreshStatus returned an error when refreshExpiry has not passed")
+	}
+	c.refreshExp = time.Now().Add(-time.Hour)
+	if err := c.refreshStatus(); err == nil {
+		t.Errorf("RefreshStatus did not return an error when refreshExpiry when it has passed")
+	}
+}
+
 func TestReusableCmd(t *testing.T) {
 	binary := "echo"
 	args := []string{"hello"}