diff --git a/GNUmakefile b/GNUmakefile index bc7b51a..51ea982 100644 --- a/GNUmakefile +++ b/GNUmakefile @@ -1,4 +1,7 @@ TEST?=$$(go list ./... |grep -v 'vendor') +# Several packages call Azure; run them one at a time so that concurrent CI jobs sharing a tenant +# do not exceed the Azure Resource Manager read limits. +TESTFLAGS?=-p 1 .PHONY: lint lint: @@ -6,17 +9,17 @@ lint: .PHONY: testrace testrace: - go test -test.v -race $(TEST) + go test $(TESTFLAGS) -test.v -race $(TEST) .PHONY: test test: - go test -test.v $(TEST) + go test $(TESTFLAGS) -test.v $(TEST) # Create a test coverage report and launch a browser to view it .PHONY: testcover testcover: if [ -f "coverage.out" ]; then rm coverage.out; fi - go test -v -coverprofile=coverage.out -covermode=count $(TEST) + go test $(TESTFLAGS) -v -coverprofile=coverage.out -covermode=count $(TEST) .PHONY: tool tool: diff --git a/alzlib.go b/alzlib.go index 5ff5a23..5f2e9fe 100644 --- a/alzlib.go +++ b/alzlib.go @@ -52,6 +52,9 @@ type AlzLib struct { cache BuiltInCache clients *azureClients mu sync.RWMutex // mu is a mutex to concurrency protect the AlzLib maps + // fetchMu serialises concurrent GetDefinitionsFromAzure callers sharing this AlzLib, so that they + // do not fetch and add the same definitions at the same time. + fetchMu sync.Mutex } type azureClients struct { @@ -65,6 +68,9 @@ type Options struct { AllowOverwrite bool // Parallelism is the number of parallel requests to make to Azure APIs when getting policy definitions // and policy set definitions. + // Requests for the same definition name are always made serially. + // Set to 1 to fetch all definitions serially. + // Requests throttled by Azure Resource Manager (HTTP 429) are retried after the delay it requests. Parallelism int // UniqueRoleDefinitions indicates whether to update the role definitions to be unique per management group. // If this is not set, you may end up with conflicting role definition names. @@ -726,7 +732,11 @@ func (az *AlzLib) Init(ctx context.Context, libs ...LibraryReference) error { // before falling back to Azure API calls. // The cache is retained for the lifetime of AlzLib; callers can explicitly clear it by calling // AddCache(nil) when they no longer need it. +// Definitions are fetched concurrently, bounded by [Options.Parallelism]. func (az *AlzLib) GetDefinitionsFromAzure(ctx context.Context, reqs []BuiltInRequest) error { + az.fetchMu.Lock() + defer az.fetchMu.Unlock() + policyDefsToGet := make([]BuiltInRequest, 0, len(reqs)) policySetDefsToGet := make([]BuiltInRequest, 0, len(reqs)) @@ -922,145 +932,177 @@ func (az *AlzLib) policySetDefinitionFromCache(name string, version *string) (*a // getBuiltInPolicies retrieves the built-in policy definitions with the given names // and adds them to the AlzLib struct. +// Requests are grouped by definition name and the groups are fetched concurrently. func (az *AlzLib) getBuiltInPolicies(ctx context.Context, reqs []BuiltInRequest) error { var ( versionedClient *armpolicy.DefinitionVersionsClient client *armpolicy.DefinitionsClient ) - for _, req := range reqs { - if az.PolicyDefinitionExists(req.ResourceID.Name, req.Version) { - continue - } - - // Check cache before calling Azure. - pdv, cacheErr := az.policyDefinitionFromCache(req.ResourceID.Name, req.Version) - if cacheErr != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicies: looking up cached built-in policy definition %s: %w", - JoinNameAndVersion(req.ResourceID.Name, req.Version), - cacheErr, - ) - } - - if pdv != nil { - if err := az.AddPolicyDefinitions(pdv); err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicies: adding cached built-in policy definition %s: %w", - JoinNameAndVersion(req.ResourceID.Name, req.Version), - err, - ) - } + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + versionedClient = policyClient.NewDefinitionVersionsClient() + client = policyClient.NewDefinitionsClient() + } - continue - } + return forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(reqs, builtInRequestName), + func(ctx context.Context, req BuiltInRequest) error { + return az.getBuiltInPolicy(ctx, client, versionedClient, req) + }, + ) +} - // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() +// getBuiltInPolicy retrieves a single built-in policy definition, from the cache if possible, +// and adds it to the AlzLib struct. +func (az *AlzLib) getBuiltInPolicy( + ctx context.Context, + client *armpolicy.DefinitionsClient, + versionedClient *armpolicy.DefinitionVersionsClient, + req BuiltInRequest, +) error { + if az.PolicyDefinitionExists(req.ResourceID.Name, req.Version) { + return nil + } - if policyClient == nil { - return errors.New("Alzlib.getBuiltInPolicies: policy client not set") - } + // Check cache before calling Azure. + pdv, cacheErr := az.policyDefinitionFromCache(req.ResourceID.Name, req.Version) + if cacheErr != nil { + return fmt.Errorf( + "Alzlib.getBuiltInPolicies: looking up cached built-in policy definition %s: %w", + JoinNameAndVersion(req.ResourceID.Name, req.Version), + cacheErr, + ) + } - // Lazily initialise clients on first Azure call. - if client == nil { - versionedClient = policyClient.NewDefinitionVersionsClient() - client = policyClient.NewDefinitionsClient() + if pdv != nil { + if err := az.AddPolicyDefinitions(pdv); err != nil { + return fmt.Errorf( + "Alzlib.getBuiltInPolicies: adding cached built-in policy definition %s: %w", + JoinNameAndVersion(req.ResourceID.Name, req.Version), + err, + ) } - var err error + return nil + } - switch req.Version { - case nil: - err = az.addLatestPolicyDefinition(ctx, client, req) - default: - err = az.addVersionedPolicyDefinition(ctx, versionedClient, req) - } + // Fall back to Azure. + if client == nil || versionedClient == nil { + return errors.New("Alzlib.getBuiltInPolicies: policy client not set") + } - if err != nil { - return err - } + if req.Version == nil { + return az.addLatestPolicyDefinition(ctx, client, req) } - return nil + return az.addVersionedPolicyDefinition(ctx, versionedClient, req) } // getBuiltInPolicySets retrieves the built-in policy set definitions with the given names // and adds them to the AlzLib struct. +// Requests are grouped by definition name and the groups are fetched concurrently. +// Once all policy set definitions have been retrieved, the policy definitions they reference +// are retrieved. func (az *AlzLib) getBuiltInPolicySets(ctx context.Context, reqs []BuiltInRequest) error { - processedRequests := make([]*assets.PolicySetDefinition, 0, len(reqs)) - var ( versionedClient *armpolicy.SetDefinitionVersionsClient client *armpolicy.SetDefinitionsClient ) - for _, req := range reqs { - if az.PolicySetDefinitionExists(req.ResourceID.Name, req.Version) { - continue - } + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + versionedClient = policyClient.NewSetDefinitionVersionsClient() + client = policyClient.NewSetDefinitionsClient() + } - // Check cache before calling Azure. - psdv, cacheErr := az.policySetDefinitionFromCache(req.ResourceID.Name, req.Version) - if cacheErr != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: looking up cached built-in policy set definition %s: %w", - req.String(), - cacheErr, - ) - } + // Each request writes only to its own slot, so the processed policy set definitions keep the + // request order regardless of the order in which they are fetched. + processed := make([]*assets.PolicySetDefinition, len(reqs)) + indices := make([]int, len(reqs)) - if psdv != nil { - if err := az.AddPolicySetDefinitions(psdv); err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: adding cached built-in policy set definition %s: %w", - req.String(), - err, - ) - } + for i := range indices { + indices[i] = i + } - processedRequests = append(processedRequests, psdv) + err := forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(indices, func(i int) string { return builtInRequestName(reqs[i]) }), + func(ctx context.Context, i int) error { + psd, err := az.getBuiltInPolicySet(ctx, client, versionedClient, reqs[i]) + processed[i] = psd - continue - } + return err + }, + ) + if err != nil { + return err + } - // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() + processedRequests := make([]*assets.PolicySetDefinition, 0, len(reqs)) - if policyClient == nil { - return errors.New("Alzlib.getBuiltInPolicySets: policy client not set") + for _, psd := range processed { + if psd != nil { + processedRequests = append(processedRequests, psd) } + } - // Lazily initialise clients on first Azure call. - if client == nil { - versionedClient = policyClient.NewSetDefinitionVersionsClient() - client = policyClient.NewSetDefinitionsClient() - } + return az.ensureReferencedPolicyDefinitions(ctx, processedRequests) +} - var ( - psd *assets.PolicySetDefinition - err error +// getBuiltInPolicySet retrieves a single built-in policy set definition, from the cache if possible, +// and adds it to the AlzLib struct. +// It returns the added policy set definition, or nil if it already existed. +func (az *AlzLib) getBuiltInPolicySet( + ctx context.Context, + client *armpolicy.SetDefinitionsClient, + versionedClient *armpolicy.SetDefinitionVersionsClient, + req BuiltInRequest, +) (*assets.PolicySetDefinition, error) { + if az.PolicySetDefinitionExists(req.ResourceID.Name, req.Version) { + return nil, nil + } + + // Check cache before calling Azure. + psdv, cacheErr := az.policySetDefinitionFromCache(req.ResourceID.Name, req.Version) + if cacheErr != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: looking up cached built-in policy set definition %s: %w", + req.String(), + cacheErr, ) + } - switch req.Version { - case nil: - psd, err = az.addLatestPolicySetDefinition(ctx, client, req) - default: - psd, err = az.addVersionedPolicySetDefinition(ctx, versionedClient, req) + if psdv != nil { + if err := az.AddPolicySetDefinitions(psdv); err != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: adding cached built-in policy set definition %s: %w", + req.String(), + err, + ) } - if err != nil { - return err - } + return psdv, nil + } - processedRequests = append(processedRequests, psd) + // Fall back to Azure. + if client == nil || versionedClient == nil { + return nil, errors.New("Alzlib.getBuiltInPolicySets: policy client not set") } - return az.ensureReferencedPolicyDefinitions(ctx, processedRequests) + if req.Version == nil { + return az.addLatestPolicySetDefinition(ctx, client, req) + } + + return az.addVersionedPolicySetDefinition(ctx, versionedClient, req) +} + +// builtInRequestName returns the definition name of a built-in request. +func builtInRequestName(req BuiltInRequest) string { + return req.ResourceID.Name } func (az *AlzLib) addLatestPolicyDefinition( @@ -1271,58 +1313,56 @@ func (az *AlzLib) ensureReferencedPolicyDefinitions( return nil } - // Lazily initialised Azure clients, shared across all references within this invocation. - var ( - definitionsClient *armpolicy.DefinitionsClient - definitionsVersionedClient *armpolicy.DefinitionVersionsClient - ) + // Validate every reference before making any Azure calls, so that malformed references fail early + // and deterministically. + refs := make([]definitionReference, 0, len(processedRequests)) for _, def := range processedRequests { for _, ref := range def.Properties.PolicyDefinitions { - err := az.ensureReferencedPolicyDefinition( - ctx, - def, - ref, - &definitionsClient, - &definitionsVersionedClient, - ) + defRef, err := parseDefinitionReference(def, ref) if err != nil { return err } + + if defRef != nil { + refs = append(refs, *defRef) + } } } - return nil -} - -func (az *AlzLib) ensureReferencedPolicyDefinition( - ctx context.Context, - setDefinition *assets.PolicySetDefinition, - ref *armpolicy.DefinitionReference, - definitionsClient **armpolicy.DefinitionsClient, - definitionsVersionedClient **armpolicy.DefinitionVersionsClient, -) error { - if ref == nil { + if len(refs) == 0 { return nil } - if ref.PolicyDefinitionID == nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: error getting policy definition references for policy set definition `%s`: "+ - "policy definition ID is nil", - JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - ) - } + var ( + definitionsClient *armpolicy.DefinitionsClient + definitionsVersionedClient *armpolicy.DefinitionVersionsClient + ) - resID, err := arm.ParseResourceID(*ref.PolicyDefinitionID) - if err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: error parsing resource id `%s` referenced in policy set `%s`", - *ref.PolicyDefinitionID, - JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - ) + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + definitionsClient = policyClient.NewDefinitionsClient() + definitionsVersionedClient = policyClient.NewDefinitionVersionsClient() } + return forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(refs, func(r definitionReference) string { return r.resID.Name }), + func(ctx context.Context, r definitionReference) error { + return az.ensureReferencedPolicyDefinition(ctx, r, definitionsClient, definitionsVersionedClient) + }, + ) +} + +func (az *AlzLib) ensureReferencedPolicyDefinition( + ctx context.Context, + defRef definitionReference, + definitionsClient *armpolicy.DefinitionsClient, + definitionsVersionedClient *armpolicy.DefinitionVersionsClient, +) error { + setDefinition, ref, resID := defRef.setDefinition, defRef.ref, defRef.resID + // Check AlzLib first to avoid redundant work. if az.PolicyDefinitionExists(resID.Name, ref.DefinitionVersion) { return nil @@ -1355,28 +1395,16 @@ func (az *AlzLib) ensureReferencedPolicyDefinition( } // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() - - if policyClient == nil { + if definitionsClient == nil || definitionsVersionedClient == nil { return errors.New("Alzlib.ensureReferencedPolicyDefinition: policy client not set") } if ref.DefinitionVersion != nil { - if *definitionsVersionedClient == nil { - *definitionsVersionedClient = policyClient.NewDefinitionVersionsClient() - } - - return az.fetchReferencedPolicyDefinitionVersions(ctx, *definitionsVersionedClient, resID.Name, setDefinition, ref) - } - - if *definitionsClient == nil { - *definitionsClient = policyClient.NewDefinitionsClient() + return az.fetchReferencedPolicyDefinitionVersions(ctx, definitionsVersionedClient, resID.Name, setDefinition, ref) } return az.fetchLatestReferencedPolicyDefinition( - ctx, *definitionsClient, resID.Name, setDefinition, ref.DefinitionVersion, + ctx, definitionsClient, resID.Name, setDefinition, ref.DefinitionVersion, ) } @@ -1424,7 +1452,7 @@ func (az *AlzLib) fetchReferencedPolicyDefinitionVersions( *ref.DefinitionVersion, definitionName, JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - az.policyDefinitions[definitionName].Versions(), + az.policyDefinitionVersionsString(definitionName), ) } diff --git a/builtin_fetch.go b/builtin_fetch.go new file mode 100644 index 0000000..216548b --- /dev/null +++ b/builtin_fetch.go @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package alzlib + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "net/http" + "strconv" + "strings" + "time" + + "github.com/Azure/alzlib/assets" + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armpolicy" + "golang.org/x/sync/errgroup" +) + +const ( + // maxThrottleRetries is the number of times a definition is fetched again after Azure Resource + // Manager throttled it (HTTP 429) beyond the retries of the Azure SDK. + maxThrottleRetries = 5 + // throttleBaseDelay is the first back-off delay when a throttled response has no Retry-After. + throttleBaseDelay = 2 * time.Second + // throttleMaxDelay caps every back-off delay, including a requested Retry-After. + throttleMaxDelay = 60 * time.Second +) + +// definitionReference is a validated policy definition reference from a built-in policy set definition. +type definitionReference struct { + setDefinition *assets.PolicySetDefinition + ref *armpolicy.DefinitionReference + resID *arm.ResourceID +} + +// effectiveParallelism returns the maximum number of concurrent Azure API calls used when fetching +// built-in definitions. Values below 1 fall back to serial fetching. +func (az *AlzLib) effectiveParallelism() int { + if az.Options == nil || az.Options.Parallelism < 1 { + return 1 + } + + return az.Options.Parallelism +} + +// policyClientFactory returns the policy client factory, or nil if it has not been set. +func (az *AlzLib) policyClientFactory() *armpolicy.ClientFactory { + az.mu.RLock() + defer az.mu.RUnlock() + + return az.clients.policyClient +} + +// policyDefinitionVersionsString returns the stored versions of the named policy definition for use in +// error messages. +func (az *AlzLib) policyDefinitionVersionsString(name string) string { + az.mu.RLock() + defer az.mu.RUnlock() + + pdvs, ok := az.policyDefinitions[name] + if !ok || pdvs == nil { + return "[]" + } + + return fmt.Sprint(pdvs.Versions()) +} + +// parseDefinitionReference validates a policy definition reference of the supplied policy set definition. +// It returns nil if the reference is nil, as there is nothing to fetch. +func parseDefinitionReference( + setDefinition *assets.PolicySetDefinition, + ref *armpolicy.DefinitionReference, +) (*definitionReference, error) { + if ref == nil { + return nil, nil + } + + if ref.PolicyDefinitionID == nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: error getting policy definition references for policy set definition `%s`: "+ + "policy definition ID is nil", + JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), + ) + } + + resID, err := arm.ParseResourceID(*ref.PolicyDefinitionID) + if err != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: error parsing resource id `%s` referenced in policy set `%s`", + *ref.PolicyDefinitionID, + JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), + ) + } + + return &definitionReference{ + setDefinition: setDefinition, + ref: ref, + resID: resID, + }, nil +} + +// groupByName groups items by case-insensitive definition name, preserving first-seen order. +// Items in one group are processed serially, so two requests for the same definition can never race +// to add it, while different definitions are fetched concurrently. +func groupByName[T any](items []T, name func(T) string) [][]T { + groups := make([][]T, 0, len(items)) + index := make(map[string]int, len(items)) + + for _, item := range items { + key := strings.ToLower(name(item)) + + if i, ok := index[key]; ok { + groups[i] = append(groups[i], item) + continue + } + + index[key] = len(groups) + groups = append(groups, []T{item}) + } + + return groups +} + +// forEachGroupConcurrently calls fn for every item, running at most limit groups at once. +// Items within a group are processed in order. +// An item throttled by Azure Resource Manager is retried, see withThrottleRetry. +// The first error cancels the remaining work and is returned. +func forEachGroupConcurrently[T any]( + ctx context.Context, + limit int, + groups [][]T, + fn func(context.Context, T) error, +) error { + if len(groups) == 0 { + return nil + } + + grp, grpCtx := errgroup.WithContext(ctx) + grp.SetLimit(max(limit, 1)) + + for _, group := range groups { + grp.Go(func() error { + for _, item := range group { + if err := grpCtx.Err(); err != nil { + return err + } + + if err := withThrottleRetry(grpCtx, func() error { return fn(grpCtx, item) }); err != nil { + return err + } + } + + return nil + }) + } + + return grp.Wait() +} + +// withThrottleRetry calls fn and, while it fails because Azure Resource Manager throttled a request +// (HTTP 429), waits and calls it again, up to maxThrottleRetries times. +// The Azure SDK already retries throttled requests, but only a few times. Concurrent fetching, and +// several processes sharing a tenant, can exhaust those retries, so a throttled definition is +// fetched again after the delay requested by Retry-After, or an exponential back-off. +// fn must be safe to call again after a failure: the fetch functions first check whether the +// definition already exists. +func withThrottleRetry(ctx context.Context, fn func() error) error { + for attempt := 0; ; attempt++ { + err := fn() + if err == nil || attempt == maxThrottleRetries || !isThrottled(err) { + return err + } + + timer := time.NewTimer(throttleDelay(err, attempt)) + + select { + case <-ctx.Done(): + timer.Stop() + + return err + case <-timer.C: + } + } +} + +// isThrottled reports whether err is an Azure Resource Manager HTTP 429 response. +func isThrottled(err error) bool { + var respErr *azcore.ResponseError + + return errors.As(err, &respErr) && respErr.StatusCode == http.StatusTooManyRequests +} + +// throttleDelay returns how long to wait before retrying after a throttled response: the delay +// requested by the response, or an exponential back-off with jitter, capped at throttleMaxDelay. +func throttleDelay(err error, attempt int) time.Duration { + var respErr *azcore.ResponseError + if errors.As(err, &respErr) && respErr.RawResponse != nil { + if d, ok := retryAfter(respErr.RawResponse.Header); ok { + return min(d, throttleMaxDelay) + } + } + + backoff := min(throttleBaseDelay<= 0 { + return time.Duration(ms) * time.Millisecond, true + } + } + } + + v := h.Get("Retry-After") + if v == "" { + return 0, false + } + + if s, err := strconv.Atoi(v); err == nil && s >= 0 { + return time.Duration(s) * time.Second, true + } + + if t, err := http.ParseTime(v); err == nil { + return max(time.Until(t), 0), true + } + + return 0, false +} diff --git a/builtin_fetch_test.go b/builtin_fetch_test.go new file mode 100644 index 0000000..fc890e1 --- /dev/null +++ b/builtin_fetch_test.go @@ -0,0 +1,449 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package alzlib + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/Azure/alzlib/internal/testutil/fakearm" + "github.com/Azure/alzlib/to" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const fakeLatency = 20 * time.Millisecond + +// newFakeAlzLib returns an AlzLib whose policy client talks to a new fake ARM server. +func newFakeAlzLib(tb testing.TB, parallelism int) (*AlzLib, *fakearm.Server) { + tb.Helper() + + srv := fakearm.New() + cf, err := srv.NewClientFactory() + require.NoError(tb, err) + + az := NewAlzLib(&Options{ + Parallelism: parallelism, + AllowOverwrite: false, + UniqueRoleDefinitions: true, + }) + az.AddPolicyClient(cf) + + return az, srv +} + +func policyDefinitionRequest(tb testing.TB, name string, version *string) BuiltInRequest { + tb.Helper() + + resID, err := arm.ParseResourceID(fmt.Sprintf("/providers/Microsoft.Authorization/policyDefinitions/%s", name)) + require.NoError(tb, err) + + return BuiltInRequest{ResourceID: resID, Version: version} +} + +func policySetDefinitionRequest(tb testing.TB, name string) BuiltInRequest { + tb.Helper() + + resID, err := arm.ParseResourceID(fmt.Sprintf("/providers/Microsoft.Authorization/policySetDefinitions/%s", name)) + require.NoError(tb, err) + + return BuiltInRequest{ResourceID: resID} +} + +// addFakePolicyDefinitions registers count policy definitions named pd-NNN and returns versionless +// requests for them. +func addFakePolicyDefinitions(tb testing.TB, srv *fakearm.Server, count int) []BuiltInRequest { + tb.Helper() + + reqs := make([]BuiltInRequest, 0, count) + + for i := range count { + name := fmt.Sprintf("pd-%03d", i) + srv.AddPolicyDefinition(name, "1.0.0") + reqs = append(reqs, policyDefinitionRequest(tb, name, nil)) + } + + return reqs +} + +func TestGetDefinitionsFromAzureRespectsParallelism(t *testing.T) { + t.Parallel() + + for _, p := range []int{1, 4, 10} { + t.Run(fmt.Sprintf("parallelism-%d", p), func(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, p) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, 40) + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, req := range reqs { + assert.True(t, az.PolicyDefinitionExists(req.ResourceID.Name, nil), req.ResourceID.Name) + assert.Equal(t, 1, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } + + assert.LessOrEqual(t, srv.PeakInFlight(), p) + + if p == 1 { + assert.Equal(t, 1, srv.PeakInFlight()) + } else { + assert.Greater(t, srv.PeakInFlight(), 1) + } + }) + } +} + +func TestGetDefinitionsFromAzureReferencedDefinitionsFetchedConcurrently(t *testing.T) { + t.Parallel() + + const parallelism = 8 + + az, srv := newFakeAlzLib(t, parallelism) + srv.Latency = fakeLatency + + refs := make([]fakearm.Reference, 0, 30) + + for i := range 30 { + name := fmt.Sprintf("ref-pd-%03d", i) + srv.AddPolicyDefinition(name, "1.0.0", "1.1.0") + + var version *string + if i%2 == 0 { + version = to.Ptr("1.*.*") + } + + refs = append(refs, fakearm.Reference{Name: name, Version: version}) + } + + srv.AddPolicySetDefinition("psd", []string{"1.0.0"}, refs...) + + reqs := []BuiltInRequest{policySetDefinitionRequest(t, "psd")} + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + require.True(t, az.PolicySetDefinitionExists("psd", nil)) + + for _, ref := range refs { + assert.True(t, az.PolicyDefinitionExists(ref.Name, ref.Version), ref.Name) + assert.Equal(t, 1, srv.Calls(ref.Name), ref.Name) + } + + assert.Greater(t, srv.PeakInFlight(), 1) + assert.LessOrEqual(t, srv.PeakInFlight(), parallelism) +} + +func TestGetDefinitionsFromAzureSharedDefinitionsFetchedOnce(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 10) + srv.Latency = fakeLatency + + shared := []string{"shared-a", "shared-b", "shared-c"} + for _, name := range shared { + srv.AddPolicyDefinition(name, "1.0.0") + } + + srv.AddPolicyDefinition("direct-dup", "1.0.0") + srv.AddPolicySetDefinition("psd-1", []string{"1.0.0"}, + fakearm.Reference{Name: "shared-a"}, + fakearm.Reference{Name: "shared-b", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "shared-c"}, + ) + srv.AddPolicySetDefinition("psd-2", []string{"1.0.0"}, + fakearm.Reference{Name: "shared-a", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "shared-b"}, + fakearm.Reference{Name: "shared-c", Version: to.Ptr("1.0.*")}, + ) + + reqs := []BuiltInRequest{ + policyDefinitionRequest(t, "direct-dup", nil), + policyDefinitionRequest(t, "direct-dup", nil), + policyDefinitionRequest(t, "direct-dup", to.Ptr("1.*.*")), + policySetDefinitionRequest(t, "psd-1"), + policySetDefinitionRequest(t, "psd-2"), + policySetDefinitionRequest(t, "psd-1"), + } + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, name := range append(shared, "direct-dup", "psd-1", "psd-2") { + assert.Equal(t, 1, srv.Calls(name), name) + } +} + +func TestGetDefinitionsFromAzureConcurrentCallers(t *testing.T) { + t.Parallel() + + const callers = 4 + + az, srv := newFakeAlzLib(t, 10) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, 20) + + var wg sync.WaitGroup + + errs := make([]error, callers) + + for i := range callers { + wg.Go(func() { + errs[i] = az.GetDefinitionsFromAzure(t.Context(), reqs) + }) + } + + wg.Wait() + + for _, err := range errs { + require.NoError(t, err) + } + + for _, req := range reqs { + assert.Equal(t, 1, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } +} + +func TestGetDefinitionsFromAzureFirstErrorCancelsRemainingWork(t *testing.T) { + t.Parallel() + + const count = 50 + + az, srv := newFakeAlzLib(t, 4) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, count) + srv.FailDefinition(reqs[0].ResourceID.Name, http.StatusNotFound) + + err := az.GetDefinitionsFromAzure(t.Context(), reqs) + require.Error(t, err) + assert.Contains(t, err.Error(), reqs[0].ResourceID.Name) + assert.Less(t, srv.TotalCalls(), count) + assert.Equal(t, 1, srv.Calls(reqs[0].ResourceID.Name), "errors other than HTTP 429 must not be retried") +} + +func TestGetDefinitionsFromAzureRetriesThrottledRequests(t *testing.T) { + t.Parallel() + + // The fake client factory disables the Azure SDK's own retries, so every retry here is alzlib's. + az, srv := newFakeAlzLib(t, 10) + reqs := addFakePolicyDefinitions(t, srv, 20) + + for _, req := range reqs { + srv.ThrottleDefinition(req.ResourceID.Name, 2, "0") + } + + srv.AddPolicyDefinition("ref-versioned", "1.0.0", "1.1.0") + srv.AddPolicyDefinition("ref-latest", "1.0.0") + srv.AddPolicySetDefinition("psd", []string{"1.0.0"}, + fakearm.Reference{Name: "ref-versioned", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "ref-latest"}, + ) + + for _, name := range []string{"psd", "ref-versioned", "ref-latest"} { + srv.ThrottleDefinition(name, maxThrottleRetries, "0") + } + + reqs = append(reqs, policySetDefinitionRequest(t, "psd")) + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, req := range reqs[:20] { + assert.True(t, az.PolicyDefinitionExists(req.ResourceID.Name, nil), req.ResourceID.Name) + assert.Equal(t, 3, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } + + assert.True(t, az.PolicySetDefinitionExists("psd", nil)) + assert.True(t, az.PolicyDefinitionExists("ref-versioned", to.Ptr("1.*.*"))) + assert.True(t, az.PolicyDefinitionExists("ref-latest", nil)) + + for _, name := range []string{"psd", "ref-versioned", "ref-latest"} { + assert.Equal(t, maxThrottleRetries+1, srv.Calls(name), name) + } +} + +func TestGetDefinitionsFromAzureGivesUpWhenThrottlingPersists(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 4) + reqs := addFakePolicyDefinitions(t, srv, 1) + srv.ThrottleDefinition(reqs[0].ResourceID.Name, -1, "0") + + err := az.GetDefinitionsFromAzure(t.Context(), reqs) + require.Error(t, err) + assert.Contains(t, err.Error(), "429") + assert.Equal(t, maxThrottleRetries+1, srv.Calls(reqs[0].ResourceID.Name)) +} + +func TestGetDefinitionsFromAzureThrottleWaitHonoursCancellation(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 4) + reqs := addFakePolicyDefinitions(t, srv, 1) + srv.ThrottleDefinition(reqs[0].ResourceID.Name, -1, "30") + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + err := az.GetDefinitionsFromAzure(ctx, reqs) + + require.Error(t, err) + assert.Less(t, time.Since(start), 5*time.Second) + assert.Equal(t, 1, srv.Calls(reqs[0].ResourceID.Name)) +} + +func TestRetryAfter(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + header http.Header + want time.Duration + ok bool + }{ + {name: "none", header: http.Header{}, ok: false}, + {name: "seconds", header: http.Header{"Retry-After": {"7"}}, want: 7 * time.Second, ok: true}, + {name: "zero", header: http.Header{"Retry-After": {"0"}}, want: 0, ok: true}, + {name: "retry-after-ms", header: http.Header{"Retry-After-Ms": {"250"}}, want: 250 * time.Millisecond, ok: true}, + {name: "x-ms-retry-after-ms", header: http.Header{"X-Ms-Retry-After-Ms": {"1500"}}, want: 1500 * time.Millisecond, ok: true}, + { + name: "milliseconds take precedence", + header: http.Header{"Retry-After": {"9"}, "Retry-After-Ms": {"100"}}, + want: 100 * time.Millisecond, + ok: true, + }, + {name: "past date", header: http.Header{"Retry-After": {"Mon, 02 Jan 2006 15:04:05 GMT"}}, want: 0, ok: true}, + {name: "invalid", header: http.Header{"Retry-After": {"soon"}}, ok: false}, + {name: "negative", header: http.Header{"Retry-After": {"-1"}}, ok: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, ok := retryAfter(tc.header) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestThrottleDelayBackoffWithoutRetryAfter(t *testing.T) { + t.Parallel() + + for attempt := range maxThrottleRetries + 1 { + backoff := min(throttleBaseDelay</[/versions]. + if len(segs) < minPathSegments || segs[0] != "providers" || segs[1] != "microsoft.authorization" { + return jsonResponse(req, http.StatusNotFound, errorBody("PathNotFound", req.URL.Path)) + } + + defType, name := segs[2], segs[3] + listVersions := len(segs) == versionsPathSegments && segs[4] == versionsSeg + + s.begin(name) + defer s.end() + + if s.Latency > 0 { + select { + case <-req.Context().Done(): + return nil, fmt.Errorf("fakearm: %w", req.Context().Err()) + case <-time.After(s.Latency): + } + } + + s.mu.Lock() + status, fail := s.failures[name] + thr, throttled := s.throttles[name] + + if throttled && thr.remaining != 0 { + thr.remaining-- + s.throttles[name] = thr + } else { + throttled = false + } + + pd, pdOK := s.policyDefinitions[name] + psd, psdOK := s.policySetDefinitions[name] + s.mu.Unlock() + + if fail { + return jsonResponse(req, status, errorBody("InjectedFailure", name)) + } + + if throttled { + resp, err := jsonResponse(req, http.StatusTooManyRequests, errorBody("TenantRequestsThrottled", name)) + if resp != nil && thr.retryAfter != "" { + resp.Header.Set("Retry-After", thr.retryAfter) + } + + return resp, err + } + + switch { + case defType == policyDefinitionsSeg && pdOK: + if listVersions { + return jsonResponse(req, http.StatusOK, policyDefinitionVersionList(pd.name, pd.versions)) + } + + return jsonResponse(req, http.StatusOK, policyDefinition(pd.name, latest(pd.versions))) + case defType == policySetDefinitionsSeg && psdOK: + if listVersions { + return jsonResponse(req, http.StatusOK, policySetDefinitionVersionList(psd)) + } + + return jsonResponse(req, http.StatusOK, policySetDefinition(psd.name, latest(psd.versions), psd.refs)) + } + + return jsonResponse(req, http.StatusNotFound, errorBody("PolicyDefinitionNotFound", name)) +} + +func (s *Server) begin(name string) { + s.mu.Lock() + defer s.mu.Unlock() + + s.calls[name]++ + s.totalCalls++ + s.inFlight++ + s.peakInFlight = max(s.peakInFlight, s.inFlight) +} + +func (s *Server) end() { + s.mu.Lock() + defer s.mu.Unlock() + + s.inFlight-- +} + +func latest(versions []string) string { + if len(versions) == 0 { + return "1.0.0" + } + + return versions[len(versions)-1] +} + +func policyDefinitionProperties(name, version string) map[string]any { + return map[string]any{ + "displayName": name, + "description": name + " description", + "policyType": "BuiltIn", + "mode": "All", + "metadata": map[string]any{}, + "parameters": map[string]any{}, + "policyRule": map[string]any{ + "if": map[string]any{"field": "type", "equals": "Microsoft.Resources/subscriptions"}, + "then": map[string]any{"effect": "audit"}, + }, + "version": version, + } +} + +func policyDefinition(name, version string) map[string]any { + return map[string]any{ + "id": fmt.Sprintf(policyDefinitionIDFmt, name), + "name": name, + "type": "Microsoft.Authorization/policyDefinitions", + "properties": policyDefinitionProperties(name, version), + } +} + +func policyDefinitionVersionList(name string, versions []string) map[string]any { + value := make([]any, 0, len(versions)) + for _, v := range versions { + value = append(value, map[string]any{ + "id": fmt.Sprintf(policyDefinitionIDFmt+"/versions/%s", name, v), + "name": v, + "type": "Microsoft.Authorization/policyDefinitions/versions", + "properties": policyDefinitionProperties(name, v), + }) + } + + return map[string]any{"value": value} +} + +func policySetDefinitionProperties(name, version string, refs []Reference) map[string]any { + pdRefs := make([]any, 0, len(refs)) + + for i, ref := range refs { + r := map[string]any{ + "policyDefinitionId": fmt.Sprintf(policyDefinitionIDFmt, ref.Name), + "policyDefinitionReferenceId": fmt.Sprintf("ref-%d", i), + "parameters": map[string]any{}, + } + if ref.Version != nil { + r["definitionVersion"] = *ref.Version + } + + pdRefs = append(pdRefs, r) + } + + return map[string]any{ + "displayName": name, + "description": name + " description", + "policyType": "BuiltIn", + "metadata": map[string]any{}, + "parameters": map[string]any{}, + "policyDefinitions": pdRefs, + "version": version, + } +} + +func policySetDefinition(name, version string, refs []Reference) map[string]any { + return map[string]any{ + "id": fmt.Sprintf(policySetDefinitionIDFmt, name), + "name": name, + "type": "Microsoft.Authorization/policySetDefinitions", + "properties": policySetDefinitionProperties(name, version, refs), + } +} + +func policySetDefinitionVersionList(psd setDefinition) map[string]any { + value := make([]any, 0, len(psd.versions)) + for _, v := range psd.versions { + value = append(value, map[string]any{ + "id": fmt.Sprintf(policySetDefinitionIDFmt+"/versions/%s", psd.name, v), + "name": v, + "type": "Microsoft.Authorization/policySetDefinitions/versions", + "properties": policySetDefinitionProperties(psd.name, v, psd.refs), + }) + } + + return map[string]any{"value": value} +} + +func errorBody(code, target string) map[string]any { + return map[string]any{ + "error": map[string]any{ + "code": code, + "message": "fakearm: " + target, + }, + } +} + +func jsonResponse(req *http.Request, status int, body any) (*http.Response, error) { + b, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("fakearm: marshalling response: %w", err) + } + + return &http.Response{ + StatusCode: status, + Status: http.StatusText(status), + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader(b)), + ContentLength: int64(len(b)), + Request: req, + }, nil +} + +// credential is a fake azcore.TokenCredential that always returns a static token. +type credential struct{} + +// GetToken implements azcore.TokenCredential. +func (credential) GetToken(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) { + return azcore.AccessToken{Token: "fake", ExpiresOn: time.Now().Add(tokenLifetime)}, nil +}