From cbabfc42062bc132b7d7c08129f6a3be13e42fd7 Mon Sep 17 00:00:00 2001
From: Jason Newman <jason@getkolla.com>
Date: Wed, 28 Aug 2024 21:19:44 +0000
Subject: [PATCH 1/3] add token source implementation

---
 kc/client.go | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/kc/client.go b/kc/client.go
index 7bb84d9..611281e 100644
--- a/kc/client.go
+++ b/kc/client.go
@@ -6,6 +6,7 @@ import (
 	"runtime/debug"
 
 	"github.com/kollalabs/sdk-go/kc/swagger"
+	"golang.org/x/oauth2"
 )
 
 const baseURL = "https://api.getkolla.com/connect"
@@ -113,6 +114,41 @@ func (c *Client) Credentials(ctx context.Context, connectorID string, consumerID
 	return creds, nil
 }
 
+type tokenSource struct {
+	connectorID string
+	consumerID  string
+
+	*Credentials
+}
+
+func (t *tokenSource) Token() (*oauth2.Token, error) {
+	if t.Credentials == nil {
+		return nil, fmt.Errorf("no credentials found")
+	}
+
+	return &oauth2.Token{
+		AccessToken: t.Credentials.Token,
+		Expiry:      t.Credentials.ExpiryTime,
+	}, nil
+}
+
+func (c *Client) CredentialsOAuth2TokenSource(ctx context.Context, connectorID string, consumerID string) (oauth2.TokenSource, error) {
+
+	// fetch the credentials
+	creds, err := c.Credentials(ctx, connectorID, consumerID)
+	if err != nil {
+		return nil, err
+	}
+
+	t := &tokenSource{
+		connectorID: connectorID,
+		consumerID:  consumerID,
+		Credentials: creds,
+	}
+
+	return t, nil
+}
+
 func setUserAgent() string {
 	userAgent := ""
 	buildInfo, ok := debug.ReadBuildInfo()

From 3592ba0cfb58dbbc78a29edf57d50e54187b9bc2 Mon Sep 17 00:00:00 2001
From: Jason Newman <jason@getkolla.com>
Date: Wed, 28 Aug 2024 21:20:46 +0000
Subject: [PATCH 2/3] keep it all internal

---
 kc/client.go | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/kc/client.go b/kc/client.go
index 611281e..e83b088 100644
--- a/kc/client.go
+++ b/kc/client.go
@@ -118,17 +118,17 @@ type tokenSource struct {
 	connectorID string
 	consumerID  string
 
-	*Credentials
+	credentials *Credentials
 }
 
 func (t *tokenSource) Token() (*oauth2.Token, error) {
-	if t.Credentials == nil {
+	if t.credentials == nil {
 		return nil, fmt.Errorf("no credentials found")
 	}
 
 	return &oauth2.Token{
-		AccessToken: t.Credentials.Token,
-		Expiry:      t.Credentials.ExpiryTime,
+		AccessToken: t.credentials.Token,
+		Expiry:      t.credentials.ExpiryTime,
 	}, nil
 }
 
@@ -143,7 +143,7 @@ func (c *Client) CredentialsOAuth2TokenSource(ctx context.Context, connectorID s
 	t := &tokenSource{
 		connectorID: connectorID,
 		consumerID:  consumerID,
-		Credentials: creds,
+		credentials: creds,
 	}
 
 	return t, nil

From 29e5c9b1be93317ca5011d63384dfa8d6a3da3fb Mon Sep 17 00:00:00 2001
From: Jason Newman <jason@getkolla.com>
Date: Wed, 28 Aug 2024 21:23:43 +0000
Subject: [PATCH 3/3] handle token refresh

---
 kc/client.go | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/kc/client.go b/kc/client.go
index e83b088..064fee0 100644
--- a/kc/client.go
+++ b/kc/client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"runtime/debug"
+	"time"
 
 	"github.com/kollalabs/sdk-go/kc/swagger"
 	"golang.org/x/oauth2"
@@ -119,31 +120,35 @@ type tokenSource struct {
 	consumerID  string
 
 	credentials *Credentials
+	kc          *Client
 }
 
 func (t *tokenSource) Token() (*oauth2.Token, error) {
-	if t.credentials == nil {
-		return nil, fmt.Errorf("no credentials found")
+	if t.kc == nil || t.connectorID == "" || t.consumerID == "" {
+		return nil, fmt.Errorf("token source not configured")
+	}
+
+	if t.credentials.ExpiryTime.Before(time.Now().Add(-time.Second * 10)) {
+		// fetch the credentials
+		creds, err := t.kc.Credentials(context.Background(), t.connectorID, t.consumerID)
+		if err != nil {
+			return nil, err
+		}
+		t.credentials = creds
 	}
 
 	return &oauth2.Token{
 		AccessToken: t.credentials.Token,
 		Expiry:      t.credentials.ExpiryTime,
 	}, nil
+
 }
 
 func (c *Client) CredentialsOAuth2TokenSource(ctx context.Context, connectorID string, consumerID string) (oauth2.TokenSource, error) {
 
-	// fetch the credentials
-	creds, err := c.Credentials(ctx, connectorID, consumerID)
-	if err != nil {
-		return nil, err
-	}
-
 	t := &tokenSource{
 		connectorID: connectorID,
 		consumerID:  consumerID,
-		credentials: creds,
 	}
 
 	return t, nil