From 9dbda0224aa9414a72eb7ae588a79c1bbbb09afd Mon Sep 17 00:00:00 2001 From: wenxuanW Date: Wed, 19 Aug 2026 21:30:36 -0700 Subject: [PATCH 1/5] Retry throttled bootstrap data requests --- pkg/bootstrapdata/bootstrap_data.go | 140 ++++++++++- pkg/bootstrapdata/bootstrap_data_test.go | 300 +++++++++++++++++++++++ 2 files changed, 434 insertions(+), 6 deletions(-) diff --git a/pkg/bootstrapdata/bootstrap_data.go b/pkg/bootstrapdata/bootstrap_data.go index f3793b6b..f727d914 100644 --- a/pkg/bootstrapdata/bootstrap_data.go +++ b/pkg/bootstrapdata/bootstrap_data.go @@ -3,14 +3,18 @@ package bootstrapdata import ( "bytes" "context" + cryptorand "crypto/rand" "encoding/json" "fmt" "io" + "math" + "math/big" "net/http" "net/url" "os" "path/filepath" "regexp" + "strconv" "strings" "time" "unicode/utf8" @@ -31,6 +35,13 @@ const ( DefaultResourceManagerEndpoint = "https://management.azure.com" DefaultAuthorityHost = "https://login.microsoftonline.com" maxResponseBytes = int64(16 << 20) + // The RP bucket refills at one request per second. Twelve hours lets a + // 30,000-node scale-out drain with headroom while keeping retries bounded. + maxThrottleRetries = 1_000 + initialThrottleRetryDelay = time.Second + maxThrottleBackoffDelay = time.Hour + bootstrapDataAttemptTimeout = 2 * time.Minute + bootstrapDataRetryTimeout = 12 * time.Hour ) var ( @@ -107,19 +118,27 @@ type Data struct { } type dependencies struct { - credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) - httpClient *http.Client + credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) + httpClient *http.Client + sleep func(context.Context, time.Duration) error + jitter func(time.Duration) time.Duration + maxThrottleRetries int + retryTimeout time.Duration } func defaultDependencies() dependencies { return dependencies{ credential: newCredential, httpClient: &http.Client{ - Timeout: 2 * time.Minute, + Timeout: bootstrapDataAttemptTimeout, CheckRedirect: func(*http.Request, []*http.Request) error { return fmt.Errorf("bootstrap-data redirects are not allowed") }, }, + sleep: sleepWithContext, + jitter: fullJitter, + maxThrottleRetries: maxThrottleRetries, + retryTimeout: bootstrapDataRetryTimeout, } } @@ -163,6 +182,12 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if audience == "" { audience = endpoint } + retryTimeout := deps.retryTimeout + if retryTimeout <= 0 { + retryTimeout = bootstrapDataRetryTimeout + } + retryCtx, cancel := context.WithTimeout(ctx, retryTimeout) + defer cancel() clientOptions := azcore.ClientOptions{Cloud: cloud.Configuration{ ActiveDirectoryAuthorityHost: options.AuthorityHost, Services: map[cloud.ServiceName]cloud.ServiceConfiguration{ @@ -173,19 +198,19 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err != nil { return nil, err } - token, err := credential.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{audience + "/.default"}}) + token, err := credential.GetToken(retryCtx, policy.TokenRequestOptions{Scopes: []string{audience + "/.default"}}) if err != nil { return nil, fmt.Errorf("acquire ARM token: %w", err) } requestURL := endpoint + options.ClusterResourceID + "/agentPools/" + options.AgentPoolName + "/listBootstrapData?api-version=" + options.APIVersion - request, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, http.NoBody) + request, err := http.NewRequestWithContext(retryCtx, http.MethodPost, requestURL, http.NoBody) if err != nil { return nil, fmt.Errorf("create bootstrap-data request: %w", err) } request.Header.Set("Authorization", "Bearer "+token.Token) request.Header.Set("Content-Type", "application/json") - response, err := deps.httpClient.Do(request) + response, err := doBootstrapDataRequest(retryCtx, request, credential, audience, deps) if err != nil { return nil, fmt.Errorf("fetch bootstrap data: %w", err) } @@ -238,6 +263,109 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro }, nil } +func doBootstrapDataRequest( + ctx context.Context, + request *http.Request, + credential azcore.TokenCredential, + audience string, + deps dependencies, +) (*http.Response, error) { + maxRetries := deps.maxThrottleRetries + if maxRetries <= 0 { + maxRetries = maxThrottleRetries + } + for retry := 0; ; retry++ { + if retry > 0 { + token, err := credential.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{audience + "/.default"}}) + if err != nil { + return nil, fmt.Errorf("refresh ARM token: %w", err) + } + request.Header.Set("Authorization", "Bearer "+token.Token) + } + response, err := deps.httpClient.Do(request.Clone(ctx)) + if err != nil { + return nil, err + } + if response.StatusCode != http.StatusTooManyRequests || retry == maxRetries { + return response, nil + } + + delay := throttleRetryDelay(response.Header.Get("Retry-After"), retry, time.Now(), deps.jitter) + if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) <= delay { + return response, nil + } + _, _ = io.Copy(io.Discard, response.Body) + _ = response.Body.Close() + sleepFn := deps.sleep + if sleepFn == nil { + sleepFn = sleepWithContext + } + if err := sleepFn(ctx, delay); err != nil { + return nil, fmt.Errorf("wait to retry bootstrap data after HTTP 429: %w", err) + } + } +} + +func throttleRetryDelay(retryAfter string, retry int, now time.Time, jitter func(time.Duration) time.Duration) time.Duration { + backoff := min(initialThrottleRetryDelay*time.Duration(1< math.MaxInt64/int64(time.Second) { + return time.Duration(math.MaxInt64), true + } + return time.Duration(seconds) * time.Second, true + } + + retryAt, err := http.ParseTime(value) + if err != nil { + return 0, false + } + return max(retryAt.Sub(now), 0), true +} + +func sleepWithContext(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + func validateOptions(options Options) error { endpoint, err := url.Parse(options.ResourceManagerEndpoint) if err != nil || endpoint.Scheme != "https" || endpoint.Host == "" || endpoint.User != nil { diff --git a/pkg/bootstrapdata/bootstrap_data_test.go b/pkg/bootstrapdata/bootstrap_data_test.go index f9d87182..ff3d049e 100644 --- a/pkg/bootstrapdata/bootstrap_data_test.go +++ b/pkg/bootstrapdata/bootstrap_data_test.go @@ -2,10 +2,12 @@ package bootstrapdata import ( "context" + "errors" "io" "net/http" "os" "path/filepath" + "strconv" "strings" "testing" "time" @@ -148,6 +150,270 @@ func TestFetchRejectsMalformedBootstrapToken(t *testing.T) { } } +func TestFetchRetriesTooManyRequests(t *testing.T) { + t.Parallel() + + const responseBody = `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + attempts := 0 + throttledBodyClosed := false + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + if attempts == 1 { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: &trackingBody{Reader: strings.NewReader("throttled"), closed: &throttledBodyClosed}, + Header: http.Header{"Retry-After": []string{"2"}}, + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(responseBody)), + Header: make(http.Header), + }, nil + })} + var delays []time.Duration + got, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + jitter: func(time.Duration) time.Duration { return 500 * time.Millisecond }, + sleep: func(_ context.Context, delay time.Duration) error { + delays = append(delays, delay) + return nil + }, + }) + if err != nil { + t.Fatalf("fetch() error = %v", err) + } + if got.BootstrapToken != "abcdef.0123456789abcdef" { + t.Fatalf("BootstrapToken = %q", got.BootstrapToken) + } + if attempts != 2 { + t.Fatalf("attempts = %d, want 2", attempts) + } + if !throttledBodyClosed { + t.Fatal("throttled response body was not closed") + } + if len(delays) != 1 || delays[0] != 2500*time.Millisecond { + t.Fatalf("retry delays = %v, want [2.5s]", delays) + } +} + +func TestFetchRefreshesARMTokenBeforeRetry(t *testing.T) { + t.Parallel() + + credential := &countingCredential{} + attempts := 0 + client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + attempts++ + wantToken := "Bearer token-" + strconv.Itoa(attempts) + if got := request.Header.Get("Authorization"); got != wantToken { + t.Fatalf("Authorization = %q, want %q", got, wantToken) + } + if attempts == 1 { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader("throttled")), + Header: http.Header{"Retry-After": []string{"1"}}, + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}`)), + Header: make(http.Header), + }, nil + })} + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) { return credential, nil }, + httpClient: client, + sleep: func(context.Context, time.Duration) error { return nil }, + }) + if err != nil { + t.Fatalf("fetch() error = %v", err) + } + if credential.calls != 2 { + t.Fatalf("GetToken calls = %d, want 2", credential.calls) + } +} + +func TestFetchStopsAfterThrottleRetriesExhausted(t *testing.T) { + t.Parallel() + + attempts := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader("throttled")), + Header: http.Header{"Retry-After": []string{"0"}}, + }, nil + })} + waits := 0 + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + maxThrottleRetries: 3, + sleep: func(context.Context, time.Duration) error { + waits++ + return nil + }, + }) + if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 429" { + t.Fatalf("fetch() error = %v, want final HTTP 429", err) + } + if attempts != 4 { + t.Fatalf("attempts = %d, want 4", attempts) + } + if waits != 3 { + t.Fatalf("waits = %d, want 3", waits) + } +} + +func TestFetchDoesNotRetryOtherStatusCodes(t *testing.T) { + t.Parallel() + + attempts := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + return &http.Response{ + StatusCode: http.StatusServiceUnavailable, + Body: io.NopCloser(strings.NewReader("unavailable")), + Header: http.Header{"Retry-After": []string{"1"}}, + }, nil + })} + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + }) + if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 503" { + t.Fatalf("fetch() error = %v, want HTTP 503", err) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) + } +} + +func TestFetchStopsRetryingWhenContextIsCancelled(t *testing.T) { + t.Parallel() + + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader("throttled")), + Header: http.Header{"Retry-After": []string{"1"}}, + }, nil + })} + ctx, cancel := context.WithCancel(t.Context()) + _, err := fetch(ctx, validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + sleep: func(ctx context.Context, _ time.Duration) error { + cancel() + return ctx.Err() + }, + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("fetch() error = %v, want context cancellation", err) + } +} + +func TestFetchDoesNotRetryBeyondDeadline(t *testing.T) { + t.Parallel() + + attempts := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader("throttled")), + Header: http.Header{"Retry-After": []string{"46800"}}, + }, nil + })} + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + sleep: func(context.Context, time.Duration) error { + t.Fatal("unexpected retry wait") + return nil + }, + }) + if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 429" { + t.Fatalf("fetch() error = %v, want final HTTP 429", err) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) + } +} + +func TestParseRetryAfter(t *testing.T) { + t.Parallel() + + now := time.Date(2026, time.August, 20, 12, 0, 0, 0, time.UTC) + tests := []struct { + name string + value string + want time.Duration + ok bool + }{ + {name: "seconds", value: "3", want: 3 * time.Second, ok: true}, + {name: "HTTP date", value: now.Add(5 * time.Second).Format(http.TimeFormat), want: 5 * time.Second, ok: true}, + {name: "past HTTP date", value: now.Add(-time.Second).Format(http.TimeFormat), want: 0, ok: true}, + {name: "negative seconds", value: "-1", ok: false}, + {name: "invalid", value: "later", ok: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, ok := parseRetryAfter(tt.value, now) + if got != tt.want || ok != tt.ok { + t.Fatalf("parseRetryAfter(%q) = (%s, %t), want (%s, %t)", tt.value, got, ok, tt.want, tt.ok) + } + }) + } +} + +func TestThrottleRetryDelayUsesServerMinimum(t *testing.T) { + t.Parallel() + + now := time.Date(2026, time.August, 20, 12, 0, 0, 0, time.UTC) + delay := throttleRetryDelay("30", 0, now, func(time.Duration) time.Duration { return 500 * time.Millisecond }) + if delay != 30500*time.Millisecond { + t.Fatalf("throttleRetryDelay() = %s, want 30.5s", delay) + } +} + +func TestThrottleRetryDelayUsesExponentialFullJitter(t *testing.T) { + t.Parallel() + + var bounds []time.Duration + for retry := range 14 { + _ = throttleRetryDelay("", retry, time.Time{}, func(bound time.Duration) time.Duration { + bounds = append(bounds, bound) + return 0 + }) + } + want := []time.Duration{ + time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, + 16 * time.Second, 32 * time.Second, 64 * time.Second, 128 * time.Second, + 256 * time.Second, 512 * time.Second, 1024 * time.Second, 2048 * time.Second, + time.Hour, time.Hour, + } + for i := range want { + if bounds[i] != want[i] { + t.Fatalf("retry %d jitter bound = %s, want %s", i, bounds[i], want[i]) + } + } +} + +func TestThrottleRetryBudgetCoversThirtyThousandNodes(t *testing.T) { + t.Parallel() + + const nodeCount = 30_000 + minimumDrainTime := nodeCount * time.Second + if bootstrapDataRetryTimeout <= minimumDrainTime { + t.Fatalf("retry timeout = %s, want more than %s", bootstrapDataRetryTimeout, minimumDrainTime) + } +} + func TestFetchAndWriteRequiresOutput(t *testing.T) { t.Parallel() @@ -164,3 +430,37 @@ func TestClientCertificateCredentialOptions(t *testing.T) { t.Fatal("SendCertificateChain = false") } } + +func validTestOptions() Options { + return Options{ + ClusterResourceID: "/subscriptions/12345678-1234-1234-1234-123456789012/resourceGroups/rg/providers/Microsoft.ContainerService/managedClusters/cluster", + AgentPoolName: "aksflexnodes", + AuthMode: "msi", + ResourceManagerEndpoint: DefaultResourceManagerEndpoint, + AuthorityHost: DefaultAuthorityHost, + APIVersion: DefaultAPIVersion, + } +} + +func staticCredentialFactory(Options, azcore.ClientOptions) (azcore.TokenCredential, error) { + return staticCredential{}, nil +} + +type trackingBody struct { + io.Reader + closed *bool +} + +type countingCredential struct { + calls int +} + +func (c *countingCredential) GetToken(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) { + c.calls++ + return azcore.AccessToken{Token: "token-" + strconv.Itoa(c.calls), ExpiresOn: time.Now().Add(time.Hour)}, nil +} + +func (b *trackingBody) Close() error { + *b.closed = true + return nil +} From 8daafe1a8d4ad174f809a2dfbcd2d319438bf6d3 Mon Sep 17 00:00:00 2001 From: wenxuanW Date: Thu, 20 Aug 2026 11:16:00 -0700 Subject: [PATCH 2/5] Use SDK for bootstrap data requests --- go.mod | 1 + go.sum | 11 +- pkg/bootstrapdata/bootstrap_data.go | 293 ++++++++++------------- pkg/bootstrapdata/bootstrap_data_test.go | 277 ++++++++------------- 4 files changed, 238 insertions(+), 344 deletions(-) diff --git a/go.mod b/go.mod index ab956db8..846686c2 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.22.0 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.14.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8 v8.3.0-beta.2 + github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9 v9.5.0-beta.1 github.com/Azure/kubelogin v0.2.15 github.com/Azure/unbounded v0.2.4 github.com/go-logr/logr v1.4.4 diff --git a/go.sum b/go.sum index 0166cd67..8265ff49 100644 --- a/go.sum +++ b/go.sum @@ -10,10 +10,15 @@ github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0 h1:fhqpLE3UEXi9lPaBRpQ6Xu github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0/go.mod h1:7dCRMLwisfRH3dBupKeNCioWYUZ4SS09Z14H+7i8ZoY= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8 v8.3.0-beta.2 h1:uTV/toeMMa4Uia3It7dRli2ePtZizYpl125iuhiH6TU= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8 v8.3.0-beta.2/go.mod h1:2lUQLQklNSBVEZfdITZzWJ84eRduPBJlM9XstZW9AWg= -github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v3 v3.1.0 h1:2qsIIvxVT+uE6yrNldntJKlLRgxGbZ85kgtz5SNBhMw= -github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v3 v3.1.0/go.mod h1:AW8VEadnhw9xox+VaVd9sP7NjzOAnaZBLRH6Tq3cJ38= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9 v9.5.0-beta.1 h1:fnsRu+aUmY9LqHMiBm+OYYwiUNp/dUHUGtjNkz0j5OY= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9 v9.5.0-beta.1/go.mod h1:VD8lsnWhQBWSA9/kT+3DI20fjha1dBy11S5XzP8dlWE= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v3 v3.2.0 h1:+lnLQhKh3cgSOIOVH61UZ3s/l9d+bAZp5d/spt1+7UI= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v3 v3.2.0/go.mod h1:tStOHrivWUrcBolspvKV70Us1ckESYGYSHdG4LX8zyY= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armdeployments v1.0.0 h1:67nFqWXpo0x5Nz0XEb1yI7s8D+EHy8NsTinYw9sZnLk= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armdeployments v1.0.0/go.mod h1:fewgRjNVE84QVVh798sIMFb7gPXPp7NmnekGnboSnXk= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0 h1:Dd+RhdJn0OTtVGaeDLZpcumkIVCtA/3/Fo42+eoYvVM= -github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0/go.mod h1:5kakwfW5CjC9KK+Q4wjXAg+ShuIm2mBMua0ZFj2C8PE= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources/v3 v3.0.1 h1:guyQA4b8XB2sbJZXzUnOF9mn0WDBv/ZT7me9wTipKtE= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources/v3 v3.0.1/go.mod h1:8h8yhzh9o+0HeSIhUxYny+rEQajScrfIpNktvgYG3Q8= github.com/Azure/go-autorest v14.2.0+incompatible h1:V5VMDjClD3GiElqLWO7mz2MxNAK/vTfRHdAubSIPRgs= github.com/Azure/go-autorest v14.2.0+incompatible/go.mod h1:r+4oMnoxhatjLLJ6zxSWATqVooLgysK6ZNox3g/xq24= github.com/Azure/go-autorest/autorest v0.11.29 h1:I4+HL/JDvErx2LjyzaVxllw2lRDB5/BT2Bm4g20iqYw= diff --git a/pkg/bootstrapdata/bootstrap_data.go b/pkg/bootstrapdata/bootstrap_data.go index f727d914..8a54d63f 100644 --- a/pkg/bootstrapdata/bootstrap_data.go +++ b/pkg/bootstrapdata/bootstrap_data.go @@ -3,18 +3,14 @@ package bootstrapdata import ( "bytes" "context" - cryptorand "crypto/rand" "encoding/json" "fmt" "io" - "math" - "math/big" "net/http" "net/url" "os" "path/filepath" "regexp" - "strconv" "strings" "time" "unicode/utf8" @@ -23,7 +19,9 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" "github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud" "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9" "github.com/google/renameio/v2" "github.com/Azure/AKSFlexNode/pkg/azclient" @@ -37,9 +35,8 @@ const ( maxResponseBytes = int64(16 << 20) // The RP bucket refills at one request per second. Twelve hours lets a // 30,000-node scale-out drain with headroom while keeping retries bounded. - maxThrottleRetries = 1_000 + maxThrottleRetries = 30_000 initialThrottleRetryDelay = time.Second - maxThrottleBackoffDelay = time.Hour bootstrapDataAttemptTimeout = 2 * time.Minute bootstrapDataRetryTimeout = 12 * time.Hour ) @@ -114,16 +111,15 @@ type Data struct { BootstrapToken string ClusterFQDN string CACertData string - raw map[string]any + raw json.RawMessage } type dependencies struct { - credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) - httpClient *http.Client - sleep func(context.Context, time.Duration) error - jitter func(time.Duration) time.Duration - maxThrottleRetries int - retryTimeout time.Duration + credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) + httpClient *http.Client + retryOptions *policy.RetryOptions + retryTimeout time.Duration + responseLimit int64 } func defaultDependencies() dependencies { @@ -135,10 +131,15 @@ func defaultDependencies() dependencies { return fmt.Errorf("bootstrap-data redirects are not allowed") }, }, - sleep: sleepWithContext, - jitter: fullJitter, - maxThrottleRetries: maxThrottleRetries, - retryTimeout: bootstrapDataRetryTimeout, + retryOptions: &policy.RetryOptions{ + MaxRetries: maxThrottleRetries, + TryTimeout: bootstrapDataAttemptTimeout, + RetryDelay: initialThrottleRetryDelay, + MaxRetryDelay: bootstrapDataRetryTimeout, + ShouldRetry: retryOnlyTooManyRequests, + }, + retryTimeout: bootstrapDataRetryTimeout, + responseLimit: maxResponseBytes, } } @@ -177,11 +178,7 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err := validateOptions(options); err != nil { return nil, err } - endpoint := strings.TrimRight(options.ResourceManagerEndpoint, "/") - audience := strings.TrimRight(options.ResourceManagerAudience, "/") - if audience == "" { - audience = endpoint - } + environment := resourceManagerEnvironment(options) retryTimeout := deps.retryTimeout if retryTimeout <= 0 { retryTimeout = bootstrapDataRetryTimeout @@ -189,181 +186,137 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro retryCtx, cancel := context.WithTimeout(ctx, retryTimeout) defer cancel() clientOptions := azcore.ClientOptions{Cloud: cloud.Configuration{ - ActiveDirectoryAuthorityHost: options.AuthorityHost, + ActiveDirectoryAuthorityHost: environment.AuthorityHost, Services: map[cloud.ServiceName]cloud.ServiceConfiguration{ - cloud.ResourceManager: {Endpoint: endpoint, Audience: audience}, + cloud.ResourceManager: {Endpoint: environment.Endpoint, Audience: environment.Audience}, }, }} credential, err := deps.credential(options, clientOptions) if err != nil { return nil, err } - token, err := credential.GetToken(retryCtx, policy.TokenRequestOptions{Scopes: []string{audience + "/.default"}}) + clusterID, err := arm.ParseResourceID(options.ClusterResourceID) if err != nil { - return nil, fmt.Errorf("acquire ARM token: %w", err) - } - requestURL := endpoint + options.ClusterResourceID + "/agentPools/" + options.AgentPoolName + - "/listBootstrapData?api-version=" + options.APIVersion - request, err := http.NewRequestWithContext(retryCtx, http.MethodPost, requestURL, http.NoBody) + return nil, fmt.Errorf("parse cluster resource ID: %w", err) + } + retryOptions := policy.RetryOptions{ + MaxRetries: maxThrottleRetries, + TryTimeout: bootstrapDataAttemptTimeout, + RetryDelay: initialThrottleRetryDelay, + MaxRetryDelay: bootstrapDataRetryTimeout, + ShouldRetry: retryOnlyTooManyRequests, + } + if deps.retryOptions != nil { + retryOptions = *deps.retryOptions + } + var transport policy.Transporter + if deps.httpClient != nil { + transport = deps.httpClient + } else { + transport = http.DefaultClient + } + responseLimit := deps.responseLimit + if responseLimit <= 0 { + responseLimit = maxResponseBytes + } + client, err := armcontainerservice.NewAgentPoolsClient(clusterID.SubscriptionID, credential, &arm.ClientOptions{ + ClientOptions: policy.ClientOptions{ + APIVersion: options.APIVersion, + Cloud: clientOptions.Cloud, + Retry: retryOptions, + Transport: limitedResponseBodyTransport{inner: transport, limit: responseLimit}, + }, + DisableRPRegistration: true, + }) if err != nil { - return nil, fmt.Errorf("create bootstrap-data request: %w", err) + return nil, fmt.Errorf("create AgentPools client: %w", err) + } + var rawResponse *http.Response + response, err := client.ListBootstrapData( + policy.WithCaptureResponse(retryCtx, &rawResponse), + clusterID.ResourceGroupName, + clusterID.Name, + options.AgentPoolName, + armcontainerservice.ListBootstrapDataRequest{}, + nil, + ) + if err != nil { + return nil, fmt.Errorf("list bootstrap data: %w", err) + } + if rawResponse == nil { + return nil, fmt.Errorf("list bootstrap data returned no HTTP response") } - request.Header.Set("Authorization", "Bearer "+token.Token) - request.Header.Set("Content-Type", "application/json") - response, err := doBootstrapDataRequest(retryCtx, request, credential, audience, deps) + raw, err := runtime.Payload(rawResponse) if err != nil { - return nil, fmt.Errorf("fetch bootstrap data: %w", err) - } - if response.StatusCode < 200 || response.StatusCode >= 300 { - _ = response.Body.Close() - return nil, fmt.Errorf("fetch bootstrap data returned HTTP status %d", response.StatusCode) - } - data, readErr := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes+1)) - closeErr := response.Body.Close() - if readErr != nil { - return nil, fmt.Errorf("read bootstrap data: %w", readErr) - } - if closeErr != nil { - return nil, fmt.Errorf("close bootstrap-data response: %w", closeErr) - } - if int64(len(data)) > maxResponseBytes { - return nil, fmt.Errorf("bootstrap data exceeds %d bytes", maxResponseBytes) - } - var raw map[string]any - if err := json.Unmarshal(data, &raw); err != nil { - return nil, fmt.Errorf("parse bootstrap data: %w", err) - } - var responseData struct { - Azure struct { - BootstrapToken struct { - Token string `json:"token"` - } `json:"bootstrapToken"` - } `json:"azure"` - Node struct { - Kubelet struct { - ClusterFQDN string `json:"clusterFQDN"` - CACertData string `json:"caCertData"` - } `json:"kubelet"` - } `json:"node"` - } - if err := json.Unmarshal(data, &responseData); err != nil { - return nil, fmt.Errorf("parse typed bootstrap data: %w", err) - } - if responseData.Azure.BootstrapToken.Token == "" { + return nil, fmt.Errorf("read bootstrap data: %w", err) + } + if int64(len(raw)) > responseLimit { + return nil, fmt.Errorf("bootstrap data exceeds %d bytes", responseLimit) + } + responseData := response.PoolBootstrapData + bootstrapToken := "" + if responseData.Azure != nil && responseData.Azure.BootstrapToken != nil && responseData.Azure.BootstrapToken.Token != nil { + bootstrapToken = *responseData.Azure.BootstrapToken.Token + } + if bootstrapToken == "" { return nil, fmt.Errorf("bootstrap-data response did not contain a bootstrap token") } - if !config.BootstrapTokenPattern.MatchString(responseData.Azure.BootstrapToken.Token) { + if !config.BootstrapTokenPattern.MatchString(bootstrapToken) { return nil, fmt.Errorf("bootstrap-data response contained an invalid bootstrap token") } - return &Data{ - BootstrapToken: responseData.Azure.BootstrapToken.Token, - ClusterFQDN: responseData.Node.Kubelet.ClusterFQDN, - CACertData: responseData.Node.Kubelet.CACertData, - raw: raw, - }, nil -} - -func doBootstrapDataRequest( - ctx context.Context, - request *http.Request, - credential azcore.TokenCredential, - audience string, - deps dependencies, -) (*http.Response, error) { - maxRetries := deps.maxThrottleRetries - if maxRetries <= 0 { - maxRetries = maxThrottleRetries - } - for retry := 0; ; retry++ { - if retry > 0 { - token, err := credential.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{audience + "/.default"}}) - if err != nil { - return nil, fmt.Errorf("refresh ARM token: %w", err) - } - request.Header.Set("Authorization", "Bearer "+token.Token) + clusterFQDN := "" + caCertData := "" + if responseData.Node != nil && responseData.Node.Kubelet != nil { + if responseData.Node.Kubelet.ClusterFQDN != nil { + clusterFQDN = *responseData.Node.Kubelet.ClusterFQDN } - response, err := deps.httpClient.Do(request.Clone(ctx)) - if err != nil { - return nil, err - } - if response.StatusCode != http.StatusTooManyRequests || retry == maxRetries { - return response, nil - } - - delay := throttleRetryDelay(response.Header.Get("Retry-After"), retry, time.Now(), deps.jitter) - if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) <= delay { - return response, nil - } - _, _ = io.Copy(io.Discard, response.Body) - _ = response.Body.Close() - sleepFn := deps.sleep - if sleepFn == nil { - sleepFn = sleepWithContext - } - if err := sleepFn(ctx, delay); err != nil { - return nil, fmt.Errorf("wait to retry bootstrap data after HTTP 429: %w", err) + if responseData.Node.Kubelet.CaCertData != nil { + caCertData = *responseData.Node.Kubelet.CaCertData } } + return &Data{ + BootstrapToken: bootstrapToken, + ClusterFQDN: clusterFQDN, + CACertData: caCertData, + raw: append(json.RawMessage(nil), raw...), + }, nil } -func throttleRetryDelay(retryAfter string, retry int, now time.Time, jitter func(time.Duration) time.Duration) time.Duration { - backoff := min(initialThrottleRetryDelay*time.Duration(1< math.MaxInt64/int64(time.Second) { - return time.Duration(math.MaxInt64), true - } - return time.Duration(seconds) * time.Second, true - } +func retryOnlyTooManyRequests(response *http.Response, err error) bool { + return err == nil && response != nil && response.StatusCode == http.StatusTooManyRequests +} - retryAt, err := http.ParseTime(value) - if err != nil { - return 0, false - } - return max(retryAt.Sub(now), 0), true +type limitedResponseBodyTransport struct { + inner policy.Transporter + limit int64 } -func sleepWithContext(ctx context.Context, delay time.Duration) error { - timer := time.NewTimer(delay) - defer timer.Stop() - select { - case <-timer.C: - return nil - case <-ctx.Done(): - return ctx.Err() - } +func (t limitedResponseBodyTransport) Do(request *http.Request) (*http.Response, error) { + response, err := t.inner.Do(request) + if err != nil || response == nil || response.Body == nil { + return response, err + } + response.Body = struct { + io.Reader + io.Closer + }{ + Reader: io.LimitReader(response.Body, t.limit+1), + Closer: response.Body, + } + return response, nil } func validateOptions(options Options) error { diff --git a/pkg/bootstrapdata/bootstrap_data_test.go b/pkg/bootstrapdata/bootstrap_data_test.go index ff3d049e..a62ca797 100644 --- a/pkg/bootstrapdata/bootstrap_data_test.go +++ b/pkg/bootstrapdata/bootstrap_data_test.go @@ -7,7 +7,6 @@ import ( "net/http" "os" "path/filepath" - "strconv" "strings" "testing" "time" @@ -42,7 +41,7 @@ func TestFetchAndWrite(t *testing.T) { if request.Header.Get("Authorization") != "Bearer arm-token" { t.Error("missing token") } - body := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + body := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}},"futureField":null}` return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}, nil })} options := Options{ @@ -65,16 +64,46 @@ func TestFetchAndWrite(t *testing.T) { if info.Mode().Perm() != 0o600 { t.Fatalf("mode = %o", info.Mode().Perm()) } + written, err := os.ReadFile(output) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(written), `"futureField": null`) { + t.Fatalf("output did not preserve unknown null field: %s", written) + } } func TestFetchInMemory(t *testing.T) { t.Parallel() + const apiVersion = "2026-06-02-preview" const response = `{ "azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}, "node":{"kubelet":{"clusterFQDN":"api.example.test","caCertData":"Y2E="}} }` - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.Method != http.MethodPost { + t.Errorf("request method = %q", request.Method) + } + if request.URL.Host != "management.usgovcloudapi.net" { + t.Errorf("request host = %q", request.URL.Host) + } + if request.URL.Path != "/subscriptions/12345678-1234-1234-1234-123456789012/resourceGroups/rg/providers/Microsoft.ContainerService/managedClusters/cluster/agentPools/aksflexnodes/listBootstrapData" { + t.Errorf("request path = %q", request.URL.Path) + } + if request.URL.Query().Get("api-version") != apiVersion { + t.Errorf("api-version = %q", request.URL.Query().Get("api-version")) + } + if request.Header.Get("Content-Type") != "application/json" { + t.Errorf("Content-Type = %q", request.Header.Get("Content-Type")) + } + requestBody, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + if string(requestBody) != "{}" { + t.Errorf("request body = %q, want {}", requestBody) + } return &http.Response{ StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(response)), @@ -88,7 +117,7 @@ func TestFetchInMemory(t *testing.T) { ResourceManagerEndpoint: "https://management.usgovcloudapi.net", ResourceManagerAudience: "https://management.core.usgovcloudapi.net", AuthorityHost: "https://login.microsoftonline.us/", - APIVersion: DefaultAPIVersion, + APIVersion: apiVersion, } var scope string got, err := fetch(t.Context(), options, dependencies{ @@ -162,7 +191,7 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { return &http.Response{ StatusCode: http.StatusTooManyRequests, Body: &trackingBody{Reader: strings.NewReader("throttled"), closed: &throttledBodyClosed}, - Header: http.Header{"Retry-After": []string{"2"}}, + Header: make(http.Header), }, nil } return &http.Response{ @@ -171,15 +200,10 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { Header: make(http.Header), }, nil })} - var delays []time.Duration got, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - jitter: func(time.Duration) time.Duration { return 500 * time.Millisecond }, - sleep: func(_ context.Context, delay time.Duration) error { - delays = append(delays, delay) - return nil - }, + credential: staticCredentialFactory, + httpClient: client, + retryOptions: noDelayRetryOptions(1), }) if err != nil { t.Fatalf("fetch() error = %v", err) @@ -193,46 +217,6 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { if !throttledBodyClosed { t.Fatal("throttled response body was not closed") } - if len(delays) != 1 || delays[0] != 2500*time.Millisecond { - t.Fatalf("retry delays = %v, want [2.5s]", delays) - } -} - -func TestFetchRefreshesARMTokenBeforeRetry(t *testing.T) { - t.Parallel() - - credential := &countingCredential{} - attempts := 0 - client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - attempts++ - wantToken := "Bearer token-" + strconv.Itoa(attempts) - if got := request.Header.Get("Authorization"); got != wantToken { - t.Fatalf("Authorization = %q, want %q", got, wantToken) - } - if attempts == 1 { - return &http.Response{ - StatusCode: http.StatusTooManyRequests, - Body: io.NopCloser(strings.NewReader("throttled")), - Header: http.Header{"Retry-After": []string{"1"}}, - }, nil - } - return &http.Response{ - StatusCode: http.StatusOK, - Body: io.NopCloser(strings.NewReader(`{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}`)), - Header: make(http.Header), - }, nil - })} - _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) { return credential, nil }, - httpClient: client, - sleep: func(context.Context, time.Duration) error { return nil }, - }) - if err != nil { - t.Fatalf("fetch() error = %v", err) - } - if credential.calls != 2 { - t.Fatalf("GetToken calls = %d, want 2", credential.calls) - } } func TestFetchStopsAfterThrottleRetriesExhausted(t *testing.T) { @@ -244,28 +228,21 @@ func TestFetchStopsAfterThrottleRetriesExhausted(t *testing.T) { return &http.Response{ StatusCode: http.StatusTooManyRequests, Body: io.NopCloser(strings.NewReader("throttled")), - Header: http.Header{"Retry-After": []string{"0"}}, + Header: make(http.Header), }, nil })} - waits := 0 _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - maxThrottleRetries: 3, - sleep: func(context.Context, time.Duration) error { - waits++ - return nil - }, + credential: staticCredentialFactory, + httpClient: client, + retryOptions: noDelayRetryOptions(3), }) - if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 429" { - t.Fatalf("fetch() error = %v, want final HTTP 429", err) + var responseError *azcore.ResponseError + if !errors.As(err, &responseError) || responseError.StatusCode != http.StatusTooManyRequests { + t.Fatalf("fetch() error = %v, want final HTTP 429 response error", err) } if attempts != 4 { t.Fatalf("attempts = %d, want 4", attempts) } - if waits != 3 { - t.Fatalf("waits = %d, want 3", waits) - } } func TestFetchDoesNotRetryOtherStatusCodes(t *testing.T) { @@ -281,136 +258,94 @@ func TestFetchDoesNotRetryOtherStatusCodes(t *testing.T) { }, nil })} _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, + credential: staticCredentialFactory, + httpClient: client, + retryOptions: noDelayRetryOptions(3), }) - if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 503" { - t.Fatalf("fetch() error = %v, want HTTP 503", err) + var responseError *azcore.ResponseError + if !errors.As(err, &responseError) || responseError.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("fetch() error = %v, want HTTP 503 response error", err) } if attempts != 1 { t.Fatalf("attempts = %d, want 1", attempts) } } -func TestFetchStopsRetryingWhenContextIsCancelled(t *testing.T) { +func TestFetchDoesNotRetryTransportErrors(t *testing.T) { t.Parallel() + attempts := 0 client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusTooManyRequests, - Body: io.NopCloser(strings.NewReader("throttled")), - Header: http.Header{"Retry-After": []string{"1"}}, - }, nil + attempts++ + return nil, errors.New("network unavailable") })} - ctx, cancel := context.WithCancel(t.Context()) - _, err := fetch(ctx, validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - sleep: func(ctx context.Context, _ time.Duration) error { - cancel() - return ctx.Err() - }, + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + retryOptions: noDelayRetryOptions(3), }) - if !errors.Is(err, context.Canceled) { - t.Fatalf("fetch() error = %v, want context cancellation", err) + if err == nil || !strings.Contains(err.Error(), "network unavailable") { + t.Fatalf("fetch() error = %v, want transport error", err) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) } } -func TestFetchDoesNotRetryBeyondDeadline(t *testing.T) { +func TestFetchStopsRetryingWhenContextIsCancelled(t *testing.T) { t.Parallel() - attempts := 0 + ctx, cancel := context.WithCancel(t.Context()) client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - attempts++ + cancel() return &http.Response{ StatusCode: http.StatusTooManyRequests, Body: io.NopCloser(strings.NewReader("throttled")), - Header: http.Header{"Retry-After": []string{"46800"}}, + Header: http.Header{"Retry-After": []string{"1"}}, }, nil })} - _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - sleep: func(context.Context, time.Duration) error { - t.Fatal("unexpected retry wait") - return nil - }, + _, err := fetch(ctx, validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + retryOptions: noDelayRetryOptions(3), }) - if err == nil || err.Error() != "fetch bootstrap data returned HTTP status 429" { - t.Fatalf("fetch() error = %v, want final HTTP 429", err) - } - if attempts != 1 { - t.Fatalf("attempts = %d, want 1", attempts) - } -} - -func TestParseRetryAfter(t *testing.T) { - t.Parallel() - - now := time.Date(2026, time.August, 20, 12, 0, 0, 0, time.UTC) - tests := []struct { - name string - value string - want time.Duration - ok bool - }{ - {name: "seconds", value: "3", want: 3 * time.Second, ok: true}, - {name: "HTTP date", value: now.Add(5 * time.Second).Format(http.TimeFormat), want: 5 * time.Second, ok: true}, - {name: "past HTTP date", value: now.Add(-time.Second).Format(http.TimeFormat), want: 0, ok: true}, - {name: "negative seconds", value: "-1", ok: false}, - {name: "invalid", value: "later", ok: false}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - got, ok := parseRetryAfter(tt.value, now) - if got != tt.want || ok != tt.ok { - t.Fatalf("parseRetryAfter(%q) = (%s, %t), want (%s, %t)", tt.value, got, ok, tt.want, tt.ok) - } - }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("fetch() error = %v, want context cancellation", err) } } -func TestThrottleRetryDelayUsesServerMinimum(t *testing.T) { +func TestThrottleRetryBudgetCoversThirtyThousandNodes(t *testing.T) { t.Parallel() - now := time.Date(2026, time.August, 20, 12, 0, 0, 0, time.UTC) - delay := throttleRetryDelay("30", 0, now, func(time.Duration) time.Duration { return 500 * time.Millisecond }) - if delay != 30500*time.Millisecond { - t.Fatalf("throttleRetryDelay() = %s, want 30.5s", delay) + const nodeCount = 30_000 + minimumDrainTime := nodeCount * time.Second + if bootstrapDataRetryTimeout <= minimumDrainTime { + t.Fatalf("retry timeout = %s, want more than %s", bootstrapDataRetryTimeout, minimumDrainTime) } -} - -func TestThrottleRetryDelayUsesExponentialFullJitter(t *testing.T) { - t.Parallel() - - var bounds []time.Duration - for retry := range 14 { - _ = throttleRetryDelay("", retry, time.Time{}, func(bound time.Duration) time.Duration { - bounds = append(bounds, bound) - return 0 - }) - } - want := []time.Duration{ - time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, - 16 * time.Second, 32 * time.Second, 64 * time.Second, 128 * time.Second, - 256 * time.Second, 512 * time.Second, 1024 * time.Second, 2048 * time.Second, - time.Hour, time.Hour, - } - for i := range want { - if bounds[i] != want[i] { - t.Fatalf("retry %d jitter bound = %s, want %s", i, bounds[i], want[i]) - } + if maxThrottleRetries < nodeCount { + t.Fatalf("max retries = %d, want at least %d", maxThrottleRetries, nodeCount) } } -func TestThrottleRetryBudgetCoversThirtyThousandNodes(t *testing.T) { +func TestFetchRejectsOversizedResponse(t *testing.T) { t.Parallel() - const nodeCount = 30_000 - minimumDrainTime := nodeCount * time.Second - if bootstrapDataRetryTimeout <= minimumDrainTime { - t.Fatalf("retry timeout = %s, want more than %s", bootstrapDataRetryTimeout, minimumDrainTime) + const responseLimit = 64 + response := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + strings.Repeat(" ", responseLimit) + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(response)), + Header: make(http.Header), + }, nil + })} + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + httpClient: client, + responseLimit: responseLimit, + }) + if err == nil || !strings.Contains(err.Error(), "bootstrap data exceeds 64 bytes") { + t.Fatalf("fetch() error = %v, want oversized response error", err) } } @@ -451,16 +386,16 @@ type trackingBody struct { closed *bool } -type countingCredential struct { - calls int -} - -func (c *countingCredential) GetToken(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) { - c.calls++ - return azcore.AccessToken{Token: "token-" + strconv.Itoa(c.calls), ExpiresOn: time.Now().Add(time.Hour)}, nil -} - func (b *trackingBody) Close() error { *b.closed = true return nil } + +func noDelayRetryOptions(maxRetries int32) *policy.RetryOptions { + return &policy.RetryOptions{ + MaxRetries: maxRetries, + RetryDelay: time.Nanosecond, + MaxRetryDelay: time.Nanosecond, + ShouldRetry: retryOnlyTooManyRequests, + } +} From bb62bac03e1aa9b5cb1c12ec2d5950a3d27112b7 Mon Sep 17 00:00:00 2001 From: wenxuanW Date: Thu, 20 Aug 2026 15:22:36 -0700 Subject: [PATCH 3/5] Simplify bootstrap data SDK client --- pkg/bootstrapdata/bootstrap_data.go | 128 ++++---------------- pkg/bootstrapdata/bootstrap_data_test.go | 147 ++--------------------- 2 files changed, 35 insertions(+), 240 deletions(-) diff --git a/pkg/bootstrapdata/bootstrap_data.go b/pkg/bootstrapdata/bootstrap_data.go index 8a54d63f..4792667f 100644 --- a/pkg/bootstrapdata/bootstrap_data.go +++ b/pkg/bootstrapdata/bootstrap_data.go @@ -5,8 +5,6 @@ import ( "context" "encoding/json" "fmt" - "io" - "net/http" "net/url" "os" "path/filepath" @@ -19,7 +17,6 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" "github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud" "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" - "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" "github.com/Azure/azure-sdk-for-go/sdk/azidentity" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9" "github.com/google/renameio/v2" @@ -32,7 +29,6 @@ const ( DefaultAPIVersion = "2026-05-02-preview" DefaultResourceManagerEndpoint = "https://management.azure.com" DefaultAuthorityHost = "https://login.microsoftonline.com" - maxResponseBytes = int64(16 << 20) // The RP bucket refills at one request per second. Twelve hours lets a // 30,000-node scale-out drain with headroom while keeping retries bounded. maxThrottleRetries = 30_000 @@ -105,42 +101,22 @@ func OptionsFromConfig(cfg *config.Config) (Options, error) { } // Data contains the short-lived Kubernetes join credentials returned by -// listBootstrapData. The raw response is retained only so the bootstrap CLI can -// write the complete RP response; runtime callers should use the typed fields. +// listBootstrapData. Runtime callers should use the typed fields. type Data struct { BootstrapToken string ClusterFQDN string CACertData string - raw json.RawMessage + raw armcontainerservice.PoolBootstrapData } type dependencies struct { - credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) - httpClient *http.Client - retryOptions *policy.RetryOptions - retryTimeout time.Duration - responseLimit int64 + credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) + transport policy.Transporter + retryOptions *policy.RetryOptions } func defaultDependencies() dependencies { - return dependencies{ - credential: newCredential, - httpClient: &http.Client{ - Timeout: bootstrapDataAttemptTimeout, - CheckRedirect: func(*http.Request, []*http.Request) error { - return fmt.Errorf("bootstrap-data redirects are not allowed") - }, - }, - retryOptions: &policy.RetryOptions{ - MaxRetries: maxThrottleRetries, - TryTimeout: bootstrapDataAttemptTimeout, - RetryDelay: initialThrottleRetryDelay, - MaxRetryDelay: bootstrapDataRetryTimeout, - ShouldRetry: retryOnlyTooManyRequests, - }, - retryTimeout: bootstrapDataRetryTimeout, - responseLimit: maxResponseBytes, - } + return dependencies{credential: newCredential} } // Fetch obtains fresh bootstrap data without persisting the sensitive response. @@ -178,17 +154,17 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err := validateOptions(options); err != nil { return nil, err } - environment := resourceManagerEnvironment(options) - retryTimeout := deps.retryTimeout - if retryTimeout <= 0 { - retryTimeout = bootstrapDataRetryTimeout - } - retryCtx, cancel := context.WithTimeout(ctx, retryTimeout) + retryCtx, cancel := context.WithTimeout(ctx, bootstrapDataRetryTimeout) defer cancel() + endpoint := strings.TrimRight(options.ResourceManagerEndpoint, "/") + audience := strings.TrimRight(options.ResourceManagerAudience, "/") + if audience == "" { + audience = endpoint + } clientOptions := azcore.ClientOptions{Cloud: cloud.Configuration{ - ActiveDirectoryAuthorityHost: environment.AuthorityHost, + ActiveDirectoryAuthorityHost: options.AuthorityHost, Services: map[cloud.ServiceName]cloud.ServiceConfiguration{ - cloud.ResourceManager: {Endpoint: environment.Endpoint, Audience: environment.Audience}, + cloud.ResourceManager: {Endpoint: endpoint, Audience: audience}, }, }} credential, err := deps.credential(options, clientOptions) @@ -199,41 +175,24 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err != nil { return nil, fmt.Errorf("parse cluster resource ID: %w", err) } - retryOptions := policy.RetryOptions{ - MaxRetries: maxThrottleRetries, - TryTimeout: bootstrapDataAttemptTimeout, - RetryDelay: initialThrottleRetryDelay, - MaxRetryDelay: bootstrapDataRetryTimeout, - ShouldRetry: retryOnlyTooManyRequests, - } + retryOptions := bootstrapDataRetryOptions() if deps.retryOptions != nil { retryOptions = *deps.retryOptions } - var transport policy.Transporter - if deps.httpClient != nil { - transport = deps.httpClient - } else { - transport = http.DefaultClient - } - responseLimit := deps.responseLimit - if responseLimit <= 0 { - responseLimit = maxResponseBytes - } client, err := armcontainerservice.NewAgentPoolsClient(clusterID.SubscriptionID, credential, &arm.ClientOptions{ ClientOptions: policy.ClientOptions{ APIVersion: options.APIVersion, Cloud: clientOptions.Cloud, Retry: retryOptions, - Transport: limitedResponseBodyTransport{inner: transport, limit: responseLimit}, + Transport: deps.transport, }, DisableRPRegistration: true, }) if err != nil { return nil, fmt.Errorf("create AgentPools client: %w", err) } - var rawResponse *http.Response response, err := client.ListBootstrapData( - policy.WithCaptureResponse(retryCtx, &rawResponse), + retryCtx, clusterID.ResourceGroupName, clusterID.Name, options.AgentPoolName, @@ -243,16 +202,6 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err != nil { return nil, fmt.Errorf("list bootstrap data: %w", err) } - if rawResponse == nil { - return nil, fmt.Errorf("list bootstrap data returned no HTTP response") - } - raw, err := runtime.Payload(rawResponse) - if err != nil { - return nil, fmt.Errorf("read bootstrap data: %w", err) - } - if int64(len(raw)) > responseLimit { - return nil, fmt.Errorf("bootstrap data exceeds %d bytes", responseLimit) - } responseData := response.PoolBootstrapData bootstrapToken := "" if responseData.Azure != nil && responseData.Azure.BootstrapToken != nil && responseData.Azure.BootstrapToken.Token != nil { @@ -278,45 +227,18 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro BootstrapToken: bootstrapToken, ClusterFQDN: clusterFQDN, CACertData: caCertData, - raw: append(json.RawMessage(nil), raw...), + raw: responseData, }, nil } -func resourceManagerEnvironment(options Options) azclient.ResourceManagerEnvironment { - cfg := &config.Config{} - cfg.Azure.ResourceManagerEndpointURL = options.ResourceManagerEndpoint - environment := azclient.ResourceManagerEnvironmentFromConfig(cfg) - if options.ResourceManagerAudience != "" { - environment.Audience = strings.TrimRight(options.ResourceManagerAudience, "/") - } - if options.AuthorityHost != "" && options.AuthorityHost != DefaultAuthorityHost { - environment.AuthorityHost = options.AuthorityHost +func bootstrapDataRetryOptions() policy.RetryOptions { + return policy.RetryOptions{ + MaxRetries: maxThrottleRetries, + TryTimeout: bootstrapDataAttemptTimeout, + RetryDelay: initialThrottleRetryDelay, + MaxRetryDelay: bootstrapDataRetryTimeout, + StatusCodes: []int{429}, } - return environment -} - -func retryOnlyTooManyRequests(response *http.Response, err error) bool { - return err == nil && response != nil && response.StatusCode == http.StatusTooManyRequests -} - -type limitedResponseBodyTransport struct { - inner policy.Transporter - limit int64 -} - -func (t limitedResponseBodyTransport) Do(request *http.Request) (*http.Response, error) { - response, err := t.inner.Do(request) - if err != nil || response == nil || response.Body == nil { - return response, err - } - response.Body = struct { - io.Reader - io.Closer - }{ - Reader: io.LimitReader(response.Body, t.limit+1), - Closer: response.Body, - } - return response, nil } func validateOptions(options Options) error { diff --git a/pkg/bootstrapdata/bootstrap_data_test.go b/pkg/bootstrapdata/bootstrap_data_test.go index a62ca797..9688c159 100644 --- a/pkg/bootstrapdata/bootstrap_data_test.go +++ b/pkg/bootstrapdata/bootstrap_data_test.go @@ -2,7 +2,6 @@ package bootstrapdata import ( "context" - "errors" "io" "net/http" "os" @@ -41,7 +40,7 @@ func TestFetchAndWrite(t *testing.T) { if request.Header.Get("Authorization") != "Bearer arm-token" { t.Error("missing token") } - body := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}},"futureField":null}` + body := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}, nil })} options := Options{ @@ -52,7 +51,7 @@ func TestFetchAndWrite(t *testing.T) { } err := fetchAndWrite(context.Background(), options, dependencies{ credential: func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) { return staticCredential{}, nil }, - httpClient: client, + transport: client, }) if err != nil { t.Fatalf("fetchAndWrite() error = %v", err) @@ -64,13 +63,6 @@ func TestFetchAndWrite(t *testing.T) { if info.Mode().Perm() != 0o600 { t.Fatalf("mode = %o", info.Mode().Perm()) } - written, err := os.ReadFile(output) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(string(written), `"futureField": null`) { - t.Fatalf("output did not preserve unknown null field: %s", written) - } } func TestFetchInMemory(t *testing.T) { @@ -128,7 +120,7 @@ func TestFetchInMemory(t *testing.T) { } return staticCredential{scope: &scope}, nil }, - httpClient: client, + transport: client, }) if err != nil { t.Fatalf("fetch() error = %v", err) @@ -169,7 +161,7 @@ func TestFetchRejectsMalformedBootstrapToken(t *testing.T) { } _, err := fetch(t.Context(), options, dependencies{ credential: func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) { return staticCredential{}, nil }, - httpClient: client, + transport: client, }) if err == nil || err.Error() != "bootstrap-data response contained an invalid bootstrap token" { t.Fatalf("fetch() error = %v, want invalid token error", err) @@ -184,13 +176,12 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { const responseBody = `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` attempts := 0 - throttledBodyClosed := false client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { attempts++ if attempts == 1 { return &http.Response{ StatusCode: http.StatusTooManyRequests, - Body: &trackingBody{Reader: strings.NewReader("throttled"), closed: &throttledBodyClosed}, + Body: io.NopCloser(strings.NewReader("throttled")), Header: make(http.Header), }, nil } @@ -202,7 +193,7 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { })} got, err := fetch(t.Context(), validTestOptions(), dependencies{ credential: staticCredentialFactory, - httpClient: client, + transport: client, retryOptions: noDelayRetryOptions(1), }) if err != nil { @@ -214,35 +205,6 @@ func TestFetchRetriesTooManyRequests(t *testing.T) { if attempts != 2 { t.Fatalf("attempts = %d, want 2", attempts) } - if !throttledBodyClosed { - t.Fatal("throttled response body was not closed") - } -} - -func TestFetchStopsAfterThrottleRetriesExhausted(t *testing.T) { - t.Parallel() - - attempts := 0 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - attempts++ - return &http.Response{ - StatusCode: http.StatusTooManyRequests, - Body: io.NopCloser(strings.NewReader("throttled")), - Header: make(http.Header), - }, nil - })} - _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - retryOptions: noDelayRetryOptions(3), - }) - var responseError *azcore.ResponseError - if !errors.As(err, &responseError) || responseError.StatusCode != http.StatusTooManyRequests { - t.Fatalf("fetch() error = %v, want final HTTP 429 response error", err) - } - if attempts != 4 { - t.Fatalf("attempts = %d, want 4", attempts) - } } func TestFetchDoesNotRetryOtherStatusCodes(t *testing.T) { @@ -259,96 +221,17 @@ func TestFetchDoesNotRetryOtherStatusCodes(t *testing.T) { })} _, err := fetch(t.Context(), validTestOptions(), dependencies{ credential: staticCredentialFactory, - httpClient: client, - retryOptions: noDelayRetryOptions(3), - }) - var responseError *azcore.ResponseError - if !errors.As(err, &responseError) || responseError.StatusCode != http.StatusServiceUnavailable { - t.Fatalf("fetch() error = %v, want HTTP 503 response error", err) - } - if attempts != 1 { - t.Fatalf("attempts = %d, want 1", attempts) - } -} - -func TestFetchDoesNotRetryTransportErrors(t *testing.T) { - t.Parallel() - - attempts := 0 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - attempts++ - return nil, errors.New("network unavailable") - })} - _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, + transport: client, retryOptions: noDelayRetryOptions(3), }) - if err == nil || !strings.Contains(err.Error(), "network unavailable") { - t.Fatalf("fetch() error = %v, want transport error", err) + if err == nil { + t.Fatal("fetch() error = nil, want HTTP 503 response error") } if attempts != 1 { t.Fatalf("attempts = %d, want 1", attempts) } } -func TestFetchStopsRetryingWhenContextIsCancelled(t *testing.T) { - t.Parallel() - - ctx, cancel := context.WithCancel(t.Context()) - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - cancel() - return &http.Response{ - StatusCode: http.StatusTooManyRequests, - Body: io.NopCloser(strings.NewReader("throttled")), - Header: http.Header{"Retry-After": []string{"1"}}, - }, nil - })} - _, err := fetch(ctx, validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - retryOptions: noDelayRetryOptions(3), - }) - if !errors.Is(err, context.Canceled) { - t.Fatalf("fetch() error = %v, want context cancellation", err) - } -} - -func TestThrottleRetryBudgetCoversThirtyThousandNodes(t *testing.T) { - t.Parallel() - - const nodeCount = 30_000 - minimumDrainTime := nodeCount * time.Second - if bootstrapDataRetryTimeout <= minimumDrainTime { - t.Fatalf("retry timeout = %s, want more than %s", bootstrapDataRetryTimeout, minimumDrainTime) - } - if maxThrottleRetries < nodeCount { - t.Fatalf("max retries = %d, want at least %d", maxThrottleRetries, nodeCount) - } -} - -func TestFetchRejectsOversizedResponse(t *testing.T) { - t.Parallel() - - const responseLimit = 64 - response := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + strings.Repeat(" ", responseLimit) - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Body: io.NopCloser(strings.NewReader(response)), - Header: make(http.Header), - }, nil - })} - _, err := fetch(t.Context(), validTestOptions(), dependencies{ - credential: staticCredentialFactory, - httpClient: client, - responseLimit: responseLimit, - }) - if err == nil || !strings.Contains(err.Error(), "bootstrap data exceeds 64 bytes") { - t.Fatalf("fetch() error = %v, want oversized response error", err) - } -} - func TestFetchAndWriteRequiresOutput(t *testing.T) { t.Parallel() @@ -381,21 +264,11 @@ func staticCredentialFactory(Options, azcore.ClientOptions) (azcore.TokenCredent return staticCredential{}, nil } -type trackingBody struct { - io.Reader - closed *bool -} - -func (b *trackingBody) Close() error { - *b.closed = true - return nil -} - func noDelayRetryOptions(maxRetries int32) *policy.RetryOptions { return &policy.RetryOptions{ MaxRetries: maxRetries, RetryDelay: time.Nanosecond, MaxRetryDelay: time.Nanosecond, - ShouldRetry: retryOnlyTooManyRequests, + StatusCodes: []int{http.StatusTooManyRequests}, } } From acbe85169f4669b9a548b1587f3837b2de53c4af Mon Sep 17 00:00:00 2001 From: wenxuanW Date: Thu, 20 Aug 2026 18:37:53 -0700 Subject: [PATCH 4/5] Harden bootstrap data retries --- pkg/bootstrapdata/bootstrap_data.go | 81 +++++++++++++++++++++--- pkg/bootstrapdata/bootstrap_data_test.go | 65 ++++++++++++++++++- 2 files changed, 135 insertions(+), 11 deletions(-) diff --git a/pkg/bootstrapdata/bootstrap_data.go b/pkg/bootstrapdata/bootstrap_data.go index 4792667f..744263cf 100644 --- a/pkg/bootstrapdata/bootstrap_data.go +++ b/pkg/bootstrapdata/bootstrap_data.go @@ -3,12 +3,16 @@ package bootstrapdata import ( "bytes" "context" + cryptorand "crypto/rand" "encoding/json" "fmt" + "math/big" + "net/http" "net/url" "os" "path/filepath" "regexp" + "strconv" "strings" "time" "unicode/utf8" @@ -17,6 +21,7 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" "github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud" "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" "github.com/Azure/azure-sdk-for-go/sdk/azidentity" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9" "github.com/google/renameio/v2" @@ -33,6 +38,8 @@ const ( // 30,000-node scale-out drain with headroom while keeping retries bounded. maxThrottleRetries = 30_000 initialThrottleRetryDelay = time.Second + initialThrottleRetryJitter = 5 * time.Minute + maxThrottleRetryJitter = time.Hour bootstrapDataAttemptTimeout = 2 * time.Minute bootstrapDataRetryTimeout = 12 * time.Hour ) @@ -106,17 +113,25 @@ type Data struct { BootstrapToken string ClusterFQDN string CACertData string - raw armcontainerservice.PoolBootstrapData + raw json.RawMessage } type dependencies struct { credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) transport policy.Transporter retryOptions *policy.RetryOptions + retryJitter func(time.Duration) time.Duration } func defaultDependencies() dependencies { - return dependencies{credential: newCredential} + return dependencies{ + credential: newCredential, + transport: &http.Client{ + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + }, + } } // Fetch obtains fresh bootstrap data without persisting the sensitive response. @@ -181,18 +196,20 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro } client, err := armcontainerservice.NewAgentPoolsClient(clusterID.SubscriptionID, credential, &arm.ClientOptions{ ClientOptions: policy.ClientOptions{ - APIVersion: options.APIVersion, - Cloud: clientOptions.Cloud, - Retry: retryOptions, - Transport: deps.transport, + APIVersion: options.APIVersion, + Cloud: clientOptions.Cloud, + Retry: retryOptions, + Transport: deps.transport, + PerRetryPolicies: []policy.Policy{&retryAfterJitterPolicy{jitter: deps.retryJitter}}, }, DisableRPRegistration: true, }) if err != nil { return nil, fmt.Errorf("create AgentPools client: %w", err) } + var rawResponse *http.Response response, err := client.ListBootstrapData( - retryCtx, + policy.WithCaptureResponse(retryCtx, &rawResponse), clusterID.ResourceGroupName, clusterID.Name, options.AgentPoolName, @@ -202,6 +219,13 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err != nil { return nil, fmt.Errorf("list bootstrap data: %w", err) } + if rawResponse == nil { + return nil, fmt.Errorf("list bootstrap data returned no HTTP response") + } + raw, err := runtime.Payload(rawResponse) + if err != nil { + return nil, fmt.Errorf("read bootstrap data: %w", err) + } responseData := response.PoolBootstrapData bootstrapToken := "" if responseData.Azure != nil && responseData.Azure.BootstrapToken != nil && responseData.Azure.BootstrapToken.Token != nil { @@ -227,7 +251,7 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro BootstrapToken: bootstrapToken, ClusterFQDN: clusterFQDN, CACertData: caCertData, - raw: responseData, + raw: append(json.RawMessage(nil), raw...), }, nil } @@ -237,8 +261,47 @@ func bootstrapDataRetryOptions() policy.RetryOptions { TryTimeout: bootstrapDataAttemptTimeout, RetryDelay: initialThrottleRetryDelay, MaxRetryDelay: bootstrapDataRetryTimeout, - StatusCodes: []int{429}, + ShouldRetry: func(response *http.Response, err error) bool { + return err == nil && response != nil && response.StatusCode == http.StatusTooManyRequests + }, + } +} + +type retryAfterJitterPolicy struct { + attempt int + jitter func(time.Duration) time.Duration +} + +func (p *retryAfterJitterPolicy) Do(request *policy.Request) (*http.Response, error) { + response, err := request.Next() + if err == nil && response != nil && response.StatusCode == http.StatusTooManyRequests { + if addRetryAfterJitter(response, p.attempt, p.jitter) { + p.attempt++ + } + } + return response, err +} + +func addRetryAfterJitter(response *http.Response, attempt int, jitter func(time.Duration) time.Duration) bool { + retryAfterSeconds, err := strconv.ParseUint(response.Header.Get("Retry-After"), 10, 31) + if err != nil || retryAfterSeconds == 0 { + return false + } + if jitter == nil { + jitter = randomRetryJitter + } + window := min(initialThrottleRetryJitter*time.Duration(1< Date: Thu, 20 Aug 2026 19:21:44 -0700 Subject: [PATCH 5/5] Limit bootstrap data responses --- pkg/bootstrapdata/bootstrap_data.go | 30 +++++++++++++++- pkg/bootstrapdata/bootstrap_data_test.go | 44 ++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/pkg/bootstrapdata/bootstrap_data.go b/pkg/bootstrapdata/bootstrap_data.go index 744263cf..93f86373 100644 --- a/pkg/bootstrapdata/bootstrap_data.go +++ b/pkg/bootstrapdata/bootstrap_data.go @@ -6,6 +6,7 @@ import ( cryptorand "crypto/rand" "encoding/json" "fmt" + "io" "math/big" "net/http" "net/url" @@ -34,6 +35,7 @@ const ( DefaultAPIVersion = "2026-05-02-preview" DefaultResourceManagerEndpoint = "https://management.azure.com" DefaultAuthorityHost = "https://login.microsoftonline.com" + maxResponseBytes = int64(16 << 20) // The RP bucket refills at one request per second. Twelve hours lets a // 30,000-node scale-out drain with headroom while keeping retries bounded. maxThrottleRetries = 30_000 @@ -194,12 +196,16 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if deps.retryOptions != nil { retryOptions = *deps.retryOptions } + transport := deps.transport + if transport == nil { + transport = http.DefaultClient + } client, err := armcontainerservice.NewAgentPoolsClient(clusterID.SubscriptionID, credential, &arm.ClientOptions{ ClientOptions: policy.ClientOptions{ APIVersion: options.APIVersion, Cloud: clientOptions.Cloud, Retry: retryOptions, - Transport: deps.transport, + Transport: limitedResponseTransport{inner: transport}, PerRetryPolicies: []policy.Policy{&retryAfterJitterPolicy{jitter: deps.retryJitter}}, }, DisableRPRegistration: true, @@ -226,6 +232,9 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err != nil { return nil, fmt.Errorf("read bootstrap data: %w", err) } + if int64(len(raw)) > maxResponseBytes { + return nil, fmt.Errorf("bootstrap data exceeds %d bytes", maxResponseBytes) + } responseData := response.PoolBootstrapData bootstrapToken := "" if responseData.Azure != nil && responseData.Azure.BootstrapToken != nil && responseData.Azure.BootstrapToken.Token != nil { @@ -304,6 +313,25 @@ func randomRetryJitter(maxDelay time.Duration) time.Duration { return time.Duration(jitter.Int64()) } +type limitedResponseTransport struct { + inner policy.Transporter +} + +func (t limitedResponseTransport) Do(request *http.Request) (*http.Response, error) { + response, err := t.inner.Do(request) + if err != nil || response == nil || response.Body == nil { + return response, err + } + response.Body = struct { + io.Reader + io.Closer + }{ + Reader: io.LimitReader(response.Body, maxResponseBytes+1), + Closer: response.Body, + } + return response, nil +} + func validateOptions(options Options) error { endpoint, err := url.Parse(options.ResourceManagerEndpoint) if err != nil || endpoint.Scheme != "https" || endpoint.Host == "" || endpoint.User != nil { diff --git a/pkg/bootstrapdata/bootstrap_data_test.go b/pkg/bootstrapdata/bootstrap_data_test.go index 04e62972..0a760355 100644 --- a/pkg/bootstrapdata/bootstrap_data_test.go +++ b/pkg/bootstrapdata/bootstrap_data_test.go @@ -291,6 +291,50 @@ func TestDefaultTransportRejectsRedirects(t *testing.T) { } } +func TestLimitedResponseTransport(t *testing.T) { + t.Parallel() + + body := strings.Repeat("x", int(maxResponseBytes)+2) + transport := limitedResponseTransport{inner: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))}, nil + })}} + request, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://example.test", http.NoBody) + if err != nil { + t.Fatal(err) + } + response, err := transport.Do(request) + if err != nil { + t.Fatal(err) + } + got, err := io.ReadAll(response.Body) + if err != nil { + t.Fatal(err) + } + if int64(len(got)) != maxResponseBytes+1 { + t.Fatalf("response size = %d, want %d", len(got), maxResponseBytes+1) + } +} + +func TestFetchRejectsOversizedResponse(t *testing.T) { + t.Parallel() + + body := `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + strings.Repeat(" ", int(maxResponseBytes)) + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + Header: make(http.Header), + }, nil + })} + _, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + transport: client, + }) + if err == nil || !strings.Contains(err.Error(), "bootstrap data exceeds") { + t.Fatalf("fetch() error = %v, want oversized response error", err) + } +} + func TestFetchAndWriteRequiresOutput(t *testing.T) { t.Parallel()