Fix test

Azure · Oct 10, 2024 · 3651576 · 3651576
1 parent 070d22e
commit 3651576
Show file tree

Hide file tree

Showing 7 changed files with 146 additions and 119 deletions.
diff --git a/common/oauthTokenManager.go b/common/oauthTokenManager.go
@@ -25,6 +25,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache"
 	"net"
 	"net/http"
 	"net/url"
@@ -41,8 +42,6 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/go-autorest/autorest/date"
 
-	"github.com/Azure/go-autorest/autorest/adal"
-
 	// importing the cache module registers the cache implementation for the current platform
 	_ "github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache"
 )
@@ -372,7 +371,7 @@ func (uotm *UserOAuthTokenManager) getTokenInfoFromEnvVar(ctx context.Context) (
 // OAuthTokenInfo contains info necessary for refresh OAuth credentials.
 type OAuthTokenInfo struct {
 	azcore.TokenCredential `json:"-"`
-	adal.Token
+	Token
 	Tenant                  string        `json:"_tenant"`
 	ActiveDirectoryEndpoint string        `json:"_ad_endpoint"`
 	LoginType               AutoLoginType `json:"_token_refresh_source"`
@@ -389,6 +388,42 @@ type OAuthTokenInfo struct {
 	Persist        bool                             `json:"_persist"`
 }
 
+// Token encapsulates the access token used to authorize Azure requests.
+// https://docs.microsoft.com/en-us/azure/active-directory/develop/v1-oauth2-client-creds-grant-flow#service-to-service-access-token-response
+type Token struct {
+	AccessToken  string `json:"access_token"`
+	RefreshToken string `json:"refresh_token"`
+
+	ExpiresIn json.Number `json:"expires_in"`
+	ExpiresOn json.Number `json:"expires_on"`
+	NotBefore json.Number `json:"not_before"`
+
+	Resource string `json:"resource"`
+	Type     string `json:"token_type"`
+}
+
+// IsZero returns true if the token object is zero-initialized.
+func (t Token) IsZero() bool {
+	return t == Token{}
+}
+
+// Expires returns the time.Time when the Token expires.
+func (t Token) Expires() time.Time {
+	s, err := t.ExpiresOn.Float64()
+	if err != nil {
+		s = -3600
+	}
+
+	expiration := date.NewUnixTimeFromSeconds(s)
+
+	return time.Time(expiration).UTC()
+}
+
+// IsExpired returns true if the Token is expired, false otherwise.
+func (t Token) IsExpired() bool {
+	return !t.Expires().After(time.Now().Add(0))
+}
+
 // IdentityInfo contains info for MSI.
 type IdentityInfo struct {
 	ClientID string `json:"_identity_client_id"`
@@ -424,7 +459,7 @@ func (identityInfo *IdentityInfo) Validate() error {
 }
 
 // Refresh gets new token with token info.
-func (credInfo *OAuthTokenInfo) Refresh(ctx context.Context) (*adal.Token, error) {
+func (credInfo *OAuthTokenInfo) Refresh(ctx context.Context) (*Token, error) {
 	// TODO: I think this method is only necessary until datalake is migrated.
 	// Returns cached TokenCredential or creates a new one if it hasn't been created yet.
 	tc, err := credInfo.GetTokenCredential()
@@ -436,7 +471,7 @@ func (credInfo *OAuthTokenInfo) Refresh(ctx context.Context) (*adal.Token, error
 	if err != nil {
 		return nil, err
 	}
-	return &adal.Token{
+	return &Token{
 		AccessToken: t.Token,
 		ExpiresOn:   json.Number(strconv.FormatInt(int64(t.ExpiresOn.Sub(date.UnixEpoch())/time.Second), 10)),
 	}, nil
@@ -646,19 +681,22 @@ func (credInfo *OAuthTokenInfo) GetDeviceCodeCredential() (azcore.TokenCredentia
 	if err != nil {
 		return nil, err
 	}
-	var persistenceOptions *azidentity.TokenCachePersistenceOptions
+	var persistentCache azidentity.Cache
 	if credInfo.Persist {
-		persistenceOptions = &azidentity.TokenCachePersistenceOptions{
+		persistentCache, err = cache.New(&cache.Options{
 			Name: TokenCache,
+		})
+		if err != nil {
+			return nil, err
 		}
 	}
-	// Read the record=
+	// Read the record
 	record := IffNotNil(credInfo.DeviceCodeInfo, azidentity.AuthenticationRecord{})
 	tc, err := azidentity.NewDeviceCodeCredential(&azidentity.DeviceCodeCredentialOptions{
 		TenantID:                       credInfo.Tenant,
 		ClientID:                       ApplicationID,
 		DisableAutomaticAuthentication: true,
-		TokenCachePersistenceOptions:   persistenceOptions,
+		Cache:                          persistentCache,
 		AuthenticationRecord:           record,
 		ClientOptions: azcore.ClientOptions{
 			Cloud:     cloud.Configuration{ActiveDirectoryAuthorityHost: authorityHost.String()},

diff --git a/common/zt_credCache_test.go b/common/zt_credCache_test.go
@@ -23,12 +23,10 @@ package common
 import (
 	"github.com/stretchr/testify/assert"
 	"testing"
-
-	"github.com/Azure/go-autorest/autorest/adal"
 )
 
 var fakeTokenInfo = OAuthTokenInfo{
-	Token: adal.Token{
+	Token: Token{
 		AccessToken:  "aaa0eXAiOiJKz1QiLCJhbGciOibbbbI1NiIsIng1dCcCImlCakwxUmNdddhpeTRmcHhJeGRacW5oTTJZayIsImtpZCI948lCakwxUmNxemhpeTRmcHhJeGRacW9oTTJZayJ9.eyJhdWQiOiJodHRwczovL3N0b3JhZ2UuYXp1cmUuY29tIiwiaXNzIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3LyIsImlhdCI6MTUyODEwNDQ5NywibmJmIjoxNTI4MTA0NDk3LCJleHAiOjE1MjgxMDgzOTcsIl9jbGFpbV9uYW1lcyI6eyJncm91aEHiOiJzcmMxIn0sIl9jbGFpbV9zb3VyY2VzIjp7InNyYzEiOnsiZW5kcG9pbnQiOiJodHRwczovL2dyYXBoLndpbmRvd3MubmV0LzcyZjk4OGJmLTg2ZjEtNDFhZi05MWFiLTJkN2NkMDExZGI0Ny91c2Vycy9hOTIzZjhkMC1kNGNlLTQyODAtOTEzNS01ZWE2ODVjMzgwMjYvZ2V0TWVtYmVyT2JqZWN0cyJ9fSwiYWNyIjoiMSIsImFpbyI6IkFVUUF1LzhIQUFBQU1nVkUzWE9DMHdQcG9OeGt1b2VsK1haVGNwOEhLekRORlp4NDZkMW5VN2VHUGNmbWdWNGxnUlN0NjUwcndXaHJPaCtaTXlGa3h2S3hVR3QvTHBjanNnPT0iLCJhbXIiOlsid2lhIiwibWZhIl0sImFwcGlkIjoiMTk1MGEyNTgtMjI3Yi00ZTMxLWE5Y2YtNzE3NDk1OTQ1ZmMyIiwiYXBwaWRhY3IiOiIwIiwiZGV2aWNlaWQiOiIyMjFjZTY3Yy1mYjM3LTQzMjYtYWJjYy0zNTRhZGJmNzk1NWYiLCJmYW1pbHlfbmFtZSI6IkZhbiIsImdpdmVuX25hbWUiOiJKYXNvbiIsImluX2NvcnAiOiJ0cnVlIiwiaXBhZGRyIjoiMTY3LjIyMC4yNTUuNTgiLCJuYW1lIjoiSmFzb24gRmFuIiwib2lkIjoiYTkyM2Y4ZDAtZDRjZS00MjgwLTkxMzUtNWVhNjg1YzM4MDI2Iiwib25wcmVtX3NpZCI6IlMtMS01LTIxLTIxNDY3NzMwODUtOTAzMzYzMjg1LTcxOTM0NDcwNy0xODI4ODgzIiwicHVpZCI6IjEwMDMwMDAwOEFCNjkzQTUi10JzY3AiOiJ1c2VyX2ltcGVyc29uYXRpb24iLCJzdWIiOiJBVVBFWXo1Y0xPd1BYcmRQaUF2OXZRamNGelpDN3dRRWd5dUJhejFfVnBFIiwidGlkIjoiNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3IiwidW5pcXVlX25hbWUiOiJqaWFjZmFuQG1pY3Jvc29mdC5jb20iLCJ1cG4iOiJqaWFjZmFuQG1pY3Jvc29mdC5jb20iLCJ1dGkiOiJfTlpKdlVQVG4wdTExTVFrTEcwTEFBIiwidmVyIjoiMS4wIn0.J3LZgQ7RTmqZzVcnsiruzLfcuK-vceNja7gp6wJhwwcPN1LzHK9Q1ANRVBKDMRulHiWvPNmavxf493EqkvgjHDkGSSTL3S7elLVF4Hr2SHHhUqyWoiEukY0jX5DT2tg71L4KujV7csJN-7ECqXyU0DSrRSRf3gCbD7c2ne5CFVCi1lEpEK_1lLiRZe45TTuJXmQrxEr4B6fY5MRkBz05lIbhxsUPmUunR02_-coNgQcHBOkdGdLGx4qjbzn58EJO0F2bimDRend3Tjnoia2aFq_kvQslcLU3BxIvYO5TZNfGkZyOlavoKEccPPmAb033zg9AKD_6_7K-R0mu1qmZUA",
 		RefreshToken: "Y2QwMTFkYjQ3LyIsImlhdCI6MTUyODEwNDQ5NywibmJmIjoxNTI4MTA0NDk3LCJleHAiOjE1MjgxMDgzOTcsIl9jbGFpbV9uYW1lcyI6eyJncm91cHMiOiJzcmMxIn0sIl9jbGFpbV9zb3VyY2VzIjp7InNyYzEiOnsiZW5kcG9pbnQiOiJodHRwczovL2dyYXBoLndpbmRvd3MubmV0LzcyZjk4OGJmLTg2ZjEtNDFhZi05MWFiLTJkN2NkMDExZGI0Ny91c2Vycy9hOTIzZjhkMC1kNGNlLTQyODAtOTEzNS01ZWE2ODVjMzgwMjYvZ2V0TWVtYmVyT2JqZWN0cyJ9fSwiYWNyIjoiMSIsImFpbyI6IkFVUUF1LzhIQUFBQU1nVkUzWE9DMHdQcG9OeGt1b2VsK1haVGNwOEhLekRORlp4NDZkMW5VN2VHUGNmbWdWNGxnUlN0NjUwcndXaHJPaCtaTXlGa3h2S3hVR3QvTHBjanNnPT0iLCJhbXIiOlsid2lhIiwibWZhIl0sImFwcGlkIjoiMTk1MGEyNTgtMjI3Yi00ZTMxLWE5Y2YtNzE3NDk1OTQ1ZmMyIiwiYXBwaWRhY3IiOiIwIiwiZGV2aWNlaWQiOiIyMjFjZTY3Yy1mYjM3LTQzMjYtYWJjYy0zNTRhZGJmNzk1NWYiLCJmYW1pbHlfbmFtZSI6IkZhbiIsImdpdmVuX25hbWUiOiJKYXNvbiIsImluX2NvcnAiOiJ0cnVlIiwiaXBhZGRyIjoiMTY3LjIyMC4yNTUuNTgiLCJuYW1lIjoiSmFzb24gRmFuIiwib2lkIjoiYTkyM2Y4ZDAtZDRjZS00MjgwLTkxMzUtNWVhNjg1YzM4MDI2Iiwib25wcmVtX3NpZCI6IlMtMS01LTIxLTIxNDY3NzMwODUtOTAzMzYzMjg1LTcxOTM0NDcwNy0xODI4ODgzIiwicHVpZCI6IjEwMDMwMDAwOEFCNjkzQTUiLCJzY3AiOiJ1c2VyX2ltcGVyc29uYXRpb24iLCJzdWIiOiJBVVBFWXo1Y0xPd1BYcmRQaUF2OXZRamNGelpDN3dRRWd5dUJhejFfVnBFIiwidGlkIjoiNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3IiwidW5pcXVlX25hbWUiOiJqaWF",
 		ExpiresIn:    "3599",
@@ -101,4 +99,4 @@ func TestCredCacheSaveLoadDeleteHas(t *testing.T) {
 	// Test has cached token, and validate remove token.
 	hasCachedToken, err = credCache.HasCachedToken()
 	a.False(hasCachedToken)
-}
+}
diff --git a/e2etest/newe2e_task_runazcopy.go b/e2etest/newe2e_task_runazcopy.go
@@ -125,6 +125,7 @@ type AzCopyEnvironment struct {
 	AzureClientId           *string `env:"AZURE_CLIENT_ID"`
 
 	InheritEnvironment bool
+	ManualLogin        bool
 }
 
 func (env *AzCopyEnvironment) generateAzcopyDir(a ScenarioAsserter) {
@@ -187,53 +188,55 @@ func (c *AzCopyCommand) applyTargetAuth(a Asserter, target ResourceManager) stri
 		// Only set it if it wasn't already configured. If it was manually configured,
 		// special testing may be occurring, and this may be indicated to just get a SAS-less URI.
 		// Alternatively, we may have already configured it here once before.
-		if c.Environment.AutoLoginMode == nil && c.Environment.ServicePrincipalAppID == nil && c.Environment.ServicePrincipalClientSecret == nil && c.Environment.AutoLoginTenantID == nil {
-			if GlobalConfig.StaticResources() {
-				c.Environment.AutoLoginMode = pointerTo("SPN")
-				oAuthInfo := GlobalConfig.E2EAuthConfig.StaticStgAcctInfo.StaticOAuth
-				a.AssertNow("At least NEW_E2E_STATIC_APPLICATION_ID and NEW_E2E_STATIC_CLIENT_SECRET must be specified to use OAuth.", Empty{true}, oAuthInfo.ApplicationID, oAuthInfo.ClientSecret)
-
-				c.Environment.ServicePrincipalAppID = &oAuthInfo.ApplicationID
-				c.Environment.ServicePrincipalClientSecret = &oAuthInfo.ClientSecret
-				c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.TenantID != "", &oAuthInfo.TenantID, nil)
-			} else {
-				// oauth should reliably work
+		if !c.Environment.ManualLogin {
+
+			if c.Environment.AutoLoginMode == nil && c.Environment.ServicePrincipalAppID == nil && c.Environment.ServicePrincipalClientSecret == nil && c.Environment.AutoLoginTenantID == nil {
+				if GlobalConfig.StaticResources() {
+					c.Environment.AutoLoginMode = pointerTo("SPN")
+					oAuthInfo := GlobalConfig.E2EAuthConfig.StaticStgAcctInfo.StaticOAuth
+					a.AssertNow("At least NEW_E2E_STATIC_APPLICATION_ID and NEW_E2E_STATIC_CLIENT_SECRET must be specified to use OAuth.", Empty{true}, oAuthInfo.ApplicationID, oAuthInfo.ClientSecret)
+
+					c.Environment.ServicePrincipalAppID = &oAuthInfo.ApplicationID
+					c.Environment.ServicePrincipalClientSecret = &oAuthInfo.ClientSecret
+					c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.TenantID != "", &oAuthInfo.TenantID, nil)
+				} else {
+					// oauth should reliably work
+					oAuthInfo := GlobalConfig.E2EAuthConfig.SubscriptionLoginInfo
+					if oAuthInfo.Environment == AzurePipeline {
+						c.Environment.InheritEnvironment = true
+						c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.DynamicOAuth.Workload.TenantId != "", &oAuthInfo.DynamicOAuth.Workload.TenantId, nil)
+						c.Environment.AutoLoginMode = pointerTo(common.EAutoLoginType.AzCLI().String())
+					} else {
+						c.Environment.AutoLoginMode = pointerTo(common.EAutoLoginType.SPN().String())
+						c.Environment.ServicePrincipalAppID = &oAuthInfo.DynamicOAuth.SPNSecret.ApplicationID
+						c.Environment.ServicePrincipalClientSecret = &oAuthInfo.DynamicOAuth.SPNSecret.ClientSecret
+						c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.DynamicOAuth.SPNSecret.TenantID != "", &oAuthInfo.DynamicOAuth.SPNSecret.TenantID, nil)
+					}
+				}
+			} else if c.Environment.AutoLoginMode != nil {
 				oAuthInfo := GlobalConfig.E2EAuthConfig.SubscriptionLoginInfo
-				if oAuthInfo.Environment == AzurePipeline {
+				if strings.ToLower(*c.Environment.AutoLoginMode) == common.EAutoLoginType.Workload().String() {
 					c.Environment.InheritEnvironment = true
-					c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.DynamicOAuth.Workload.TenantId != "", &oAuthInfo.DynamicOAuth.Workload.TenantId, nil)
-					c.Environment.AutoLoginMode = pointerTo(common.EAutoLoginType.AzCLI().String())
-				} else {
-					c.Environment.AutoLoginMode = pointerTo(common.EAutoLoginType.SPN().String())
-					c.Environment.ServicePrincipalAppID = &oAuthInfo.DynamicOAuth.SPNSecret.ApplicationID
-					c.Environment.ServicePrincipalClientSecret = &oAuthInfo.DynamicOAuth.SPNSecret.ClientSecret
-					c.Environment.AutoLoginTenantID = common.Iff(oAuthInfo.DynamicOAuth.SPNSecret.TenantID != "", &oAuthInfo.DynamicOAuth.SPNSecret.TenantID, nil)
+					// Get the value of the AZURE_FEDERATED_TOKEN environment variable
+					token := oAuthInfo.DynamicOAuth.Workload.FederatedToken
+					a.AssertNow("idToken must be specified to authenticate with workload identity", Empty{Invert: true}, token)
+					// Write the token to a temporary file
+					// Create a temporary file to store the token
+					file, err := os.CreateTemp("", "azure_federated_token.txt")
+					a.AssertNow("Error creating temporary file", IsNil{}, err)
+					defer file.Close()
+
+					// Write the token to the temporary file
+					_, err = file.WriteString(token)
+					a.AssertNow("Error writing to temporary file", IsNil{}, err)
+
+					// Set the AZURE_FEDERATED_TOKEN_FILE environment variable
+					c.Environment.AzureFederatedTokenFile = pointerTo(file.Name())
+					c.Environment.AzureTenantId = pointerTo(oAuthInfo.DynamicOAuth.Workload.TenantId)
+					c.Environment.AzureClientId = pointerTo(oAuthInfo.DynamicOAuth.Workload.ClientId)
 				}
 			}
-		} else if c.Environment.AutoLoginMode != nil {
-			oAuthInfo := GlobalConfig.E2EAuthConfig.SubscriptionLoginInfo
-			if strings.ToLower(*c.Environment.AutoLoginMode) == common.EAutoLoginType.Workload().String() {
-				c.Environment.InheritEnvironment = true
-				// Get the value of the AZURE_FEDERATED_TOKEN environment variable
-				token := oAuthInfo.DynamicOAuth.Workload.FederatedToken
-				a.AssertNow("idToken must be specified to authenticate with workload identity", Empty{Invert: true}, token)
-				// Write the token to a temporary file
-				// Create a temporary file to store the token
-				file, err := os.CreateTemp("", "azure_federated_token.txt")
-				a.AssertNow("Error creating temporary file", IsNil{}, err)
-				defer file.Close()
-
-				// Write the token to the temporary file
-				_, err = file.WriteString(token)
-				a.AssertNow("Error writing to temporary file", IsNil{}, err)
-
-				// Set the AZURE_FEDERATED_TOKEN_FILE environment variable
-				c.Environment.AzureFederatedTokenFile = pointerTo(file.Name())
-				c.Environment.AzureTenantId = pointerTo(oAuthInfo.DynamicOAuth.Workload.TenantId)
-				c.Environment.AzureClientId = pointerTo(oAuthInfo.DynamicOAuth.Workload.ClientId)
-			}
 		}
-
 		return target.URI(opts) // Generate like public
 	default:
 		a.Error("unsupported credential type")
@@ -339,8 +342,10 @@ func RunAzCopy(a ScenarioAsserter, commandSpec AzCopyCommand) (AzCopyStdout, *Az
 		0, command.ProcessState.ExitCode())
 
 	a.Cleanup(func(a ScenarioAsserter) {
-		UploadLogs(a, out, stderr, DerefOrZero(commandSpec.Environment.LogLocation))
-		_ = os.RemoveAll(DerefOrZero(commandSpec.Environment.LogLocation))
+		if !commandSpec.Environment.ManualLogin {
+			UploadLogs(a, out, stderr, DerefOrZero(commandSpec.Environment.LogLocation))
+			_ = os.RemoveAll(DerefOrZero(commandSpec.Environment.LogLocation))
+		}
 	})
 
 	return out, &AzCopyJobPlan{}

diff --git a/e2etest/newe2e_task_runazcopy_parameters.go b/e2etest/newe2e_task_runazcopy_parameters.go
@@ -20,7 +20,7 @@ import (
 // - "defaultfunc:SpecialDefault" Calls a `func Default*(a ScenarioAsserter) string` named SpecialDefault on the struct. Asserts if not found.
 // - "serializer:SerializerFunc" Calls a `func Serialize*(value any, a ScenarioAsserter) string` named SerializerFunc on the struct. Asserts if not found.
 // If special characters , or : are for some reason used, \ can be used as an escape.
-func MapFromTags(val reflect.Value, tagName string, a ScenarioAsserter) map[string]string {
+func MapFromTags(val reflect.Value, tagName string, a Asserter) map[string]string {
 	queue := []reflect.Value{val}
 	out := make(map[string]string)