Expired CLI authentication tokens will now automatically be refreshed f…

…ixes #261 (#234)

azure/service.go

@@ -10,10 +10,12 @@ import (
 	"regexp"
 	"runtime"
 	"strings"
+	"time"
 
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/Azure/go-autorest/autorest/azure/auth"
+	"github.com/Azure/go-autorest/autorest/azure/cli"
 	"github.com/turbot/steampipe-plugin-sdk/plugin"
 )
 
@@ -22,6 +24,7 @@ type Session struct {
 	SubscriptionID string
 	TenantID       string
 	Authorizer     autorest.Authorizer
+	Expires        *time.Time
 }
 
 // GetNewSession creates an session configured from environment variables/CLI in the order:
@@ -31,6 +34,16 @@ type Session struct {
 // 4. MSI
 // 5. CLI
 func GetNewSession(ctx context.Context, d *plugin.QueryData, tokenAudience string) (session *Session, err error) {
+	cacheKey := "GetNewSession"
+	if cachedData, ok := d.ConnectionManager.Cache.Get(cacheKey); ok {
+		session = cachedData.(*Session)
+		if session.Expires != nil && WillExpireIn(*session.Expires, 0) {
+			d.ConnectionManager.Cache.Delete("GetNewSession")
+		} else {
+			return cachedData.(*Session), nil
+		}
+	}
+
 	azureConfig := GetConfig(d.Connection)
 
 	if azureConfig.TenantID != nil {
@@ -69,13 +82,7 @@ func GetNewSession(ctx context.Context, d *plugin.QueryData, tokenAudience strin
 	}
 
 	var authorizer autorest.Authorizer
-
-	// have we already created and cached the session?
-	serviceCacheKey := tokenAudience + resource + authMethod
-
-	if cachedData, ok := d.ConnectionManager.Cache.Get(serviceCacheKey); ok {
-		return cachedData.(*Session), nil
-	}
+	var expiresOn time.Time
 
 	// so if it was not in cache - create session
 	switch authMethod {
@@ -100,7 +107,16 @@ func GetNewSession(ctx context.Context, d *plugin.QueryData, tokenAudience strin
 			return nil, err
 		}
 	default:
-		authorizer, err = auth.NewAuthorizerFromCLIWithResource(resource)
+		// authorizer, err = auth.NewAuthorizerFromCLIWithResource(resource)
+		token, err := cli.GetTokenFromCLI(resource)
+		if err != nil {
+			return nil, err
+		}
+
+		// var adalToken adal.Token
+		adalToken, err := token.ToADALToken()
+		expiresOn = adalToken.Expires()
+
 		if err != nil {
 			logger.Debug("GetNewSession__", "NewAuthorizerFromCLIWithResource error", err)
 
@@ -109,6 +125,7 @@ func GetNewSession(ctx context.Context, d *plugin.QueryData, tokenAudience strin
 			}
 			return nil, err
 		}
+		authorizer = autorest.NewBearerAuthorizer(&adalToken)
 	}
 
 	if authMethod == "CLI" {
@@ -129,9 +146,14 @@ func GetNewSession(ctx context.Context, d *plugin.QueryData, tokenAudience strin
 		SubscriptionID: subscriptionID,
 		Authorizer:     authorizer,
 		TenantID:       tenantID,
+		Expires:        &expiresOn,
 	}
 
-	d.ConnectionManager.Cache.Set(serviceCacheKey, sess)
+	if sess.Expires != nil {
+		d.ConnectionManager.Cache.SetWithTTL(cacheKey, sess, time.Until(*sess.Expires))
+	} else {
+		d.ConnectionManager.Cache.Set(cacheKey, sess)
+	}
 
 	return sess, err
 }
@@ -253,3 +275,9 @@ func getSubscriptionFromCLI(resource string) (*subscription, error) {
 		TenantID:       tokenResponse["tenant"].(string),
 	}, nil
 }
+
+// WillExpireIn returns true if the Token will expire after the passed time.Duration interval
+// from now, false otherwise.
+func WillExpireIn(t time.Time, d time.Duration) bool {
+	return !t.After(time.Now().Add(d))
+}

go.mod

@@ -7,9 +7,10 @@ require (
 	github.com/Azure/azure-storage-blob-go v0.12.0
 	github.com/Azure/go-autorest/autorest v0.11.17
 	github.com/Azure/go-autorest/autorest/azure/auth v0.5.6
+	github.com/Azure/go-autorest/autorest/azure/cli v0.4.2
 	github.com/Azure/go-autorest/autorest/date v0.3.0
 	github.com/gofrs/uuid v4.0.0+incompatible // indirect
 	github.com/tombuildsstuff/giovanni v0.15.1
 	github.com/turbot/go-kit v0.2.2-0.20210628165333-268ba0a30be3
-	github.com/turbot/steampipe-plugin-sdk v1.4.1
+	github.com/turbot/steampipe-plugin-sdk v1.5.0
 )

go.sum

@@ -120,6 +120,8 @@ github.com/hashicorp/go-multierror v1.0.0 h1:iVjPR7a6H0tWELX5NxNe7bYopibicUzc7uP
 github.com/hashicorp/go-multierror v1.0.0/go.mod h1:dHtQlpGsu+cZNNAkkCN/P3hoUDHhCYQXV3UM06sGGrk=
 github.com/hashicorp/go-plugin v1.4.1 h1:6UltRQlLN9iZO513VveELp5xyaFxVD2+1OVylE+2E+w=
 github.com/hashicorp/go-plugin v1.4.1/go.mod h1:5fGEH17QVwTTcR0zV7yhDPLLmFX9YSZ38b18Udy6vYQ=
+github.com/hashicorp/go-version v1.3.0 h1:McDWVJIU/y+u1BRV06dPaLfLCaT7fUTJLp5r04x7iNw=
+github.com/hashicorp/go-version v1.3.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
 github.com/hashicorp/hcl/v2 v2.9.1 h1:eOy4gREY0/ZQHNItlfuEZqtcQbXIxzojlP301hDpnac=
 github.com/hashicorp/hcl/v2 v2.9.1/go.mod h1:FwWsfWEjyV/CMj8s/gqAuiviY72rJ1/oayI9WftqcKg=
 github.com/hashicorp/yamux v0.0.0-20180604194846-3520598351bb h1:b5rjCoWHc7eqmAS4/qyk21ZsHyb6Mxv/jykxvNTkU4M=
@@ -183,8 +185,8 @@ github.com/tombuildsstuff/giovanni v0.15.1 h1:CVRaLOJ7C/eercCrKIsarfJ4SZoGMdBL9Q
 github.com/tombuildsstuff/giovanni v0.15.1/go.mod h1:0TZugJPEtqzPlMpuJHYfXY6Dq2uLPrXf98D2XQSxNbA=
 github.com/turbot/go-kit v0.2.2-0.20210628165333-268ba0a30be3 h1:UAfWYp+K7oESlqomRus4k+h/dSPXU17tEcarbRdtBwQ=
 github.com/turbot/go-kit v0.2.2-0.20210628165333-268ba0a30be3/go.mod h1:SBdPRngbEfYubiR81iAVtO43oPkg1+ASr+XxvgbH7/k=
-github.com/turbot/steampipe-plugin-sdk v1.4.1 h1:kY9A2JWvb3dKEzpgVK3ctTJKL/hELpcO6FaaXokmKns=
-github.com/turbot/steampipe-plugin-sdk v1.4.1/go.mod h1:j15bT2BWdzmuBuueULfoHLtSPFz1Yer9xyDvCpggetM=
+github.com/turbot/steampipe-plugin-sdk v1.5.0 h1:LX3pq307RAQc57cj7qD/EN0P7a7vQFznbM4bbnfRiKc=
+github.com/turbot/steampipe-plugin-sdk v1.5.0/go.mod h1:zM68yGM+wjkjDPz8yUTx6078GDTVWkrI26EC56XIspw=
 github.com/vmihailenco/msgpack v3.3.3+incompatible/go.mod h1:fy3FlTQTDXWkZ7Bh6AcGMlsjHatGryHQYUTf1ShIgkk=
 github.com/vmihailenco/msgpack/v4 v4.3.12/go.mod h1:gborTTJjAo/GWTqqRjrLCn9pgNN+NXzzngzBKDPIqw4=
 github.com/vmihailenco/tagparser v0.1.1/go.mod h1:OeAg3pn3UbLjkWt+rN9oFYB6u/cQgqMEUPoW2WPyhdI=