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 f3793b6b..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" @@ -20,6 +18,7 @@ import ( "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/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v9" "github.com/google/renameio/v2" "github.com/Azure/AKSFlexNode/pkg/azclient" @@ -30,7 +29,12 @@ 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 + initialThrottleRetryDelay = time.Second + bootstrapDataAttemptTimeout = 2 * time.Minute + bootstrapDataRetryTimeout = 12 * time.Hour ) var ( @@ -97,30 +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 map[string]any + raw armcontainerservice.PoolBootstrapData } type dependencies struct { - credential func(Options, azcore.ClientOptions) (azcore.TokenCredential, error) - httpClient *http.Client + 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: 2 * time.Minute, - CheckRedirect: func(*http.Request, []*http.Request) error { - return fmt.Errorf("bootstrap-data redirects are not allowed") - }, - }, - } + return dependencies{credential: newCredential} } // Fetch obtains fresh bootstrap data without persisting the sensitive response. @@ -158,6 +154,8 @@ func fetch(ctx context.Context, options Options, deps dependencies) (*Data, erro if err := validateOptions(options); err != nil { return nil, err } + retryCtx, cancel := context.WithTimeout(ctx, bootstrapDataRetryTimeout) + defer cancel() endpoint := strings.TrimRight(options.ResourceManagerEndpoint, "/") audience := strings.TrimRight(options.ResourceManagerAudience, "/") if audience == "" { @@ -173,71 +171,76 @@ 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"}}) + 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(ctx, http.MethodPost, requestURL, http.NoBody) + return nil, fmt.Errorf("parse cluster resource ID: %w", err) + } + retryOptions := bootstrapDataRetryOptions() + if deps.retryOptions != nil { + retryOptions = *deps.retryOptions + } + client, err := armcontainerservice.NewAgentPoolsClient(clusterID.SubscriptionID, credential, &arm.ClientOptions{ + ClientOptions: policy.ClientOptions{ + APIVersion: options.APIVersion, + Cloud: clientOptions.Cloud, + Retry: retryOptions, + Transport: deps.transport, + }, + DisableRPRegistration: true, + }) 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) + return nil, fmt.Errorf("create AgentPools client: %w", err) + } + response, err := client.ListBootstrapData( + retryCtx, + clusterID.ResourceGroupName, + clusterID.Name, + options.AgentPoolName, + armcontainerservice.ListBootstrapDataRequest{}, + nil, + ) 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("list bootstrap data: %w", err) + } + 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") } + clusterFQDN := "" + caCertData := "" + if responseData.Node != nil && responseData.Node.Kubelet != nil { + if responseData.Node.Kubelet.ClusterFQDN != nil { + clusterFQDN = *responseData.Node.Kubelet.ClusterFQDN + } + if responseData.Node.Kubelet.CaCertData != nil { + caCertData = *responseData.Node.Kubelet.CaCertData + } + } return &Data{ - BootstrapToken: responseData.Azure.BootstrapToken.Token, - ClusterFQDN: responseData.Node.Kubelet.ClusterFQDN, - CACertData: responseData.Node.Kubelet.CACertData, - raw: raw, + BootstrapToken: bootstrapToken, + ClusterFQDN: clusterFQDN, + CACertData: caCertData, + raw: responseData, }, nil } +func bootstrapDataRetryOptions() policy.RetryOptions { + return policy.RetryOptions{ + MaxRetries: maxThrottleRetries, + TryTimeout: bootstrapDataAttemptTimeout, + RetryDelay: initialThrottleRetryDelay, + MaxRetryDelay: bootstrapDataRetryTimeout, + StatusCodes: []int{429}, + } +} + 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..9688c159 100644 --- a/pkg/bootstrapdata/bootstrap_data_test.go +++ b/pkg/bootstrapdata/bootstrap_data_test.go @@ -51,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) @@ -68,11 +68,34 @@ func TestFetchAndWrite(t *testing.T) { 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)), @@ -86,7 +109,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{ @@ -97,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) @@ -138,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) @@ -148,6 +171,67 @@ func TestFetchRejectsMalformedBootstrapToken(t *testing.T) { } } +func TestFetchRetriesTooManyRequests(t *testing.T) { + t.Parallel() + + const responseBody = `{"azure":{"bootstrapToken":{"token":"abcdef.0123456789abcdef"}}}` + attempts := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + if attempts == 1 { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader("throttled")), + Header: make(http.Header), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(responseBody)), + Header: make(http.Header), + }, nil + })} + got, err := fetch(t.Context(), validTestOptions(), dependencies{ + credential: staticCredentialFactory, + transport: client, + retryOptions: noDelayRetryOptions(1), + }) + 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) + } +} + +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, + transport: client, + retryOptions: noDelayRetryOptions(3), + }) + if err == nil { + t.Fatal("fetch() error = nil, want HTTP 503 response error") + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) + } +} + func TestFetchAndWriteRequiresOutput(t *testing.T) { t.Parallel() @@ -164,3 +248,27 @@ 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 +} + +func noDelayRetryOptions(maxRetries int32) *policy.RetryOptions { + return &policy.RetryOptions{ + MaxRetries: maxRetries, + RetryDelay: time.Nanosecond, + MaxRetryDelay: time.Nanosecond, + StatusCodes: []int{http.StatusTooManyRequests}, + } +}