Merge pull request #68 from with-joy/fix/67-auto-token-refresh

fix(auth): implement automatic OAuth token refresh

Author: platinummonkey

pkg/client/client.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/DataDog/datadog-api-client-go/v2/api/datadog"
+	"github.com/DataDog/pup/pkg/auth/dcr"
 	"github.com/DataDog/pup/pkg/auth/storage"
 	"github.com/DataDog/pup/pkg/config"
 	"github.com/DataDog/pup/pkg/useragent"
@@ -25,6 +26,12 @@ type Client struct {
 	api    *datadog.APIClient
 }
 
+// Test hooks — overridden in tests to inject fakes
+var (
+	getStorageFunc   = func() (storage.Storage, error) { return storage.GetStorage(nil) }
+	newDCRClientFunc = func(site string) *dcr.Client { return dcr.NewClient(site) }
+)
+
 // New creates a new Datadog API client
 // Authentication priority:
 //  1. OAuth2 tokens (if available and valid)
@@ -45,16 +52,31 @@ func NewWithOptions(cfg *config.Config, forceAPIKeys bool) (*Client, error) {
 
 	if !forceAPIKeys {
 		// Try OAuth2 tokens first (preferred method)
-		store, err := storage.GetStorage(nil)
+		store, err := getStorageFunc()
 		if err == nil {
 			tokens, err := store.LoadTokens(cfg.Site)
-			if err == nil && tokens != nil && !tokens.IsExpired() {
-				// Use OAuth2 Bearer token authentication
-				ctx = context.WithValue(
-					context.Background(),
-					datadog.ContextAccessToken,
-					tokens.AccessToken,
-				)
+			if err == nil && tokens != nil {
+				// Auto-refresh: if token is expired but refresh token is available, refresh it
+				if tokens.IsExpired() && tokens.RefreshToken != "" {
+					creds, credsErr := store.LoadClientCredentials(cfg.Site)
+					if credsErr == nil && creds != nil {
+						dcrClient := newDCRClientFunc(cfg.Site)
+						newTokens, refreshErr := dcrClient.RefreshToken(tokens.RefreshToken, creds)
+						if refreshErr == nil {
+							_ = store.SaveTokens(cfg.Site, newTokens)
+							tokens = newTokens
+						}
+					}
+				}
+
+				if !tokens.IsExpired() {
+					// Use OAuth2 Bearer token authentication
+					ctx = context.WithValue(
+						context.Background(),
+						datadog.ContextAccessToken,
+						tokens.AccessToken,
+					)
+				}
 			}
 		}
 	}