No public description PiperOrigin-RevId: 958597527
diff --git a/cmd/mwlid_extension/refresher.go b/cmd/mwlid_extension/refresher.go index 8ffa443..cee8de0 100644 --- a/cmd/mwlid_extension/refresher.go +++ b/cmd/mwlid_extension/refresher.go
@@ -30,6 +30,7 @@ "github.com/GoogleCloudPlatform/galog" "github.com/GoogleCloudPlatform/google-guest-agent/internal/cfg" "github.com/GoogleCloudPlatform/google-guest-agent/internal/metadata" + "github.com/GoogleCloudPlatform/google-guest-agent/internal/retry" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" @@ -105,6 +106,8 @@ // grpcClient is the client connection to the grpc server. This is used to // cache the connection so it doesn't need to be recreated every time. grpcClient *grpc.ClientConn + // retryPolicy is the retry policy for gRPC calls. + retryPolicy retry.Policy } // outputOpts is a struct for output directory name and symlink templates. @@ -473,6 +476,39 @@ j.grpcClient = nil } +// shouldRetryGRPC checks if the gRPC error is retriable. +func shouldRetryGRPC(err error) bool { + switch status.Code(err) { + // gRPC equivalents for HTTP 500 internal server errors. Terminate retry loop. + case codes.Internal, codes.DataLoss, codes.Unknown: + return false + // Permanent client or precondition errors. Terminate retry loop. + case codes.FailedPrecondition, codes.InvalidArgument, + codes.Unauthenticated, codes.PermissionDenied, codes.Unimplemented: + return false + // On gRPC equivalents for HTTP 503 (codes.Unavailable) and other transient errors, retry. + default: + return true + } +} + +func (j *RefresherJob) getGRPCRetryPolicy() retry.Policy { + p := j.retryPolicy + if p.MaxAttempts == 0 { + p.MaxAttempts = 5 + } + if p.BackoffFactor == 0 { + p.BackoffFactor = 1 + } + if p.Jitter == 0 { + p.Jitter = time.Second + } + if p.ShouldRetry == nil { + p.ShouldRetry = shouldRetryGRPC + } + return p +} + // refreshCredsWithGRPC refreshes the workload certificates using the MWLID // service over gRPC. func (j *RefresherJob) refreshCredsWithGRPC(ctx context.Context, contentDir string) error { @@ -488,17 +524,23 @@ // Close the client if we failed to make a successful call to the server. Next // attempt to connect to the server will recreate the client. - tCtxCerts, cancelCerts := context.WithTimeout(ctx, defaultGRPCTimeout) - defer cancelCerts() - certs, err := c.GetWorkloadCertificates(tCtxCerts, &wipb.GetWorkloadCertificatesRequest{}) + policy := j.getGRPCRetryPolicy() + + certs, err := retry.RunWithResponse(ctx, policy, func() (*wipb.GetWorkloadCertificatesResponse, error) { + tCtxCerts, cancelCerts := context.WithTimeout(ctx, defaultGRPCTimeout) + defer cancelCerts() + return c.GetWorkloadCertificates(tCtxCerts, &wipb.GetWorkloadCertificatesRequest{}) + }) if err != nil { j.closeClient() return fmt.Errorf("failed to get workload certificates: %w", err) } - tCtxBundle, cancelBundle := context.WithTimeout(ctx, defaultGRPCTimeout) - defer cancelBundle() - bundle, err := c.GetWorkloadTrustBundles(tCtxBundle, &wipb.GetWorkloadTrustBundlesRequest{}) + bundle, err := retry.RunWithResponse(ctx, policy, func() (*wipb.GetWorkloadTrustBundlesResponse, error) { + tCtxBundle, cancelBundle := context.WithTimeout(ctx, defaultGRPCTimeout) + defer cancelBundle() + return c.GetWorkloadTrustBundles(tCtxBundle, &wipb.GetWorkloadTrustBundlesRequest{}) + }) if err != nil { j.closeClient() return fmt.Errorf("failed to get workload trust bundles: %w", err)
diff --git a/cmd/mwlid_extension/refresher_test.go b/cmd/mwlid_extension/refresher_test.go index 4e98ec2..331366e 100644 --- a/cmd/mwlid_extension/refresher_test.go +++ b/cmd/mwlid_extension/refresher_test.go
@@ -30,6 +30,7 @@ "github.com/GoogleCloudPlatform/google-guest-agent/internal/cfg" "github.com/GoogleCloudPlatform/google-guest-agent/internal/metadata" + "github.com/GoogleCloudPlatform/google-guest-agent/internal/retry" "github.com/GoogleCloudPlatform/google-guest-agent/internal/utils/file" "github.com/google/go-cmp/cmp" "google.golang.org/grpc" @@ -567,21 +568,33 @@ wipb.UnimplementedWorkloadIdentityServer getWorkloadCertificatesErr error getWorkloadTrustBundlesErr error + getWorkloadCertificatesErrFn func(reqCount int) error + getWorkloadTrustBundlesErrFn func(reqCount int) error reqCount int + trustBundlesReqCount int getWorkloadCertificatesResponse *wipb.GetWorkloadCertificatesResponse getWorkloadTrustBundlesResponse *wipb.GetWorkloadTrustBundlesResponse } func (s *mockWorkloadIdentityServer) GetWorkloadCertificates(ctx context.Context, req *wipb.GetWorkloadCertificatesRequest) (*wipb.GetWorkloadCertificatesResponse, error) { s.reqCount++ - if s.getWorkloadCertificatesErr != nil { + if s.getWorkloadCertificatesErrFn != nil { + if err := s.getWorkloadCertificatesErrFn(s.reqCount); err != nil { + return nil, err + } + } else if s.getWorkloadCertificatesErr != nil { return nil, s.getWorkloadCertificatesErr } return s.getWorkloadCertificatesResponse, nil } func (s *mockWorkloadIdentityServer) GetWorkloadTrustBundles(ctx context.Context, req *wipb.GetWorkloadTrustBundlesRequest) (*wipb.GetWorkloadTrustBundlesResponse, error) { - if s.getWorkloadTrustBundlesErr != nil { + s.trustBundlesReqCount++ + if s.getWorkloadTrustBundlesErrFn != nil { + if err := s.getWorkloadTrustBundlesErrFn(s.trustBundlesReqCount); err != nil { + return nil, err + } + } else if s.getWorkloadTrustBundlesErr != nil { return nil, s.getWorkloadTrustBundlesErr } return s.getWorkloadTrustBundlesResponse, nil @@ -826,6 +839,283 @@ } } +func TestShouldRetryGRPC(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "Unavailable_503_Retries", + err: status.Error(codes.Unavailable, "service unavailable"), + want: true, + }, + { + name: "Internal_500_Terminates", + err: status.Error(codes.Internal, "internal error"), + want: false, + }, + { + name: "DataLoss_500_Terminates", + err: status.Error(codes.DataLoss, "data loss error"), + want: false, + }, + { + name: "Unknown_500_Terminates", + err: status.Error(codes.Unknown, "unknown error"), + want: false, + }, + { + name: "InvalidArgument_Terminates", + err: status.Error(codes.InvalidArgument, "invalid argument"), + want: false, + }, + { + name: "NonGRPCError_Terminates", + err: fmt.Errorf("generic error"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := shouldRetryGRPC(tt.err); got != tt.want { + t.Errorf("shouldRetryGRPC(%v) = %v, want %v", tt.err, got, tt.want) + } + }) + } +} + +func TestGetGRPCRetryPolicy(t *testing.T) { + tests := []struct { + name string + initial retry.Policy + wantMax int + wantBackoff float64 + wantJitter time.Duration + testErr error + wantRetryRes bool + }{ + { + name: "DefaultsApplied", + initial: retry.Policy{}, + wantMax: 5, + wantBackoff: 1, + wantJitter: time.Second, + testErr: status.Error(codes.Unavailable, "503 unavailable"), + wantRetryRes: true, + }, + { + name: "ExistingValuesPreserved", + initial: retry.Policy{ + MaxAttempts: 10, + BackoffFactor: 2.5, + Jitter: time.Minute, + ShouldRetry: func(err error) bool { return false }, + }, + wantMax: 10, + wantBackoff: 2.5, + wantJitter: time.Minute, + testErr: status.Error(codes.Unavailable, "503 unavailable"), + wantRetryRes: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + job := &RefresherJob{retryPolicy: tt.initial} + got := job.getGRPCRetryPolicy() + + if got.MaxAttempts != tt.wantMax { + t.Errorf("getGRPCRetryPolicy().MaxAttempts = %d, want %d", got.MaxAttempts, tt.wantMax) + } + if got.BackoffFactor != tt.wantBackoff { + t.Errorf("getGRPCRetryPolicy().BackoffFactor = %v, want %v", got.BackoffFactor, tt.wantBackoff) + } + if got.Jitter != tt.wantJitter { + t.Errorf("getGRPCRetryPolicy().Jitter = %v, want %v", got.Jitter, tt.wantJitter) + } + if got.ShouldRetry == nil { + t.Fatalf("getGRPCRetryPolicy().ShouldRetry is nil") + } + if res := got.ShouldRetry(tt.testErr); res != tt.wantRetryRes { + t.Errorf("getGRPCRetryPolicy().ShouldRetry(%v) = %v, want %v", tt.testErr, res, tt.wantRetryRes) + } + }) + } +} + +func TestRefreshCredsWithGRPCRetry503(t *testing.T) { + ctx := context.Background() + if err := cfg.Load(nil); err != nil { + t.Fatalf("cfg.Load() failed unexpectedly with error: %v", err) + } + config := cfg.Retrieve() + config.MWLID.Enabled = true + + mockServer := &mockWorkloadIdentityServer{ + getWorkloadCertificatesErr: status.Error(codes.Unavailable, "service unavailable (503)"), + } + + addr, stop := startTestGRPCServer(t, mockServer) + defer stop() + + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + t.Fatalf("failed to parse address: %v", err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("failed to parse port: %v", err) + } + + config.MWLID.ServiceIP = host + config.MWLID.ServicePort = port + + job := &RefresherJob{ + retryPolicy: retry.Policy{ + MaxAttempts: 3, + Jitter: time.Millisecond, + }, + } + job.setStatus(ServiceAvailable) + + tmpDir := t.TempDir() + opts := outputOpts{ + contentDirPrefix: filepath.Join(tmpDir, "contents"), + tempSymlinkPrefix: filepath.Join(tmpDir, "symlink"), + symlink: filepath.Join(tmpDir, "creds"), + } + + err = job.refreshCreds(ctx, opts, time.Now().Format(time.RFC3339)) + if err == nil { + t.Errorf("refreshCreds expected error for 503 unavailable, got nil") + } + + if mockServer.reqCount != 3 { + t.Errorf("mockServer.reqCount = %d, want 3 (retried 503 error)", mockServer.reqCount) + } +} + +func TestRefreshCredsWithGRPCTerminate500(t *testing.T) { + ctx := context.Background() + if err := cfg.Load(nil); err != nil { + t.Fatalf("cfg.Load() failed unexpectedly with error: %v", err) + } + config := cfg.Retrieve() + config.MWLID.Enabled = true + + mockServer := &mockWorkloadIdentityServer{ + getWorkloadCertificatesErr: status.Error(codes.Internal, "internal server error (500)"), + } + + addr, stop := startTestGRPCServer(t, mockServer) + defer stop() + + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + t.Fatalf("failed to parse address: %v", err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("failed to parse port: %v", err) + } + + config.MWLID.ServiceIP = host + config.MWLID.ServicePort = port + + job := &RefresherJob{ + retryPolicy: retry.Policy{ + MaxAttempts: 3, + Jitter: time.Millisecond, + }, + } + job.setStatus(ServiceAvailable) + + tmpDir := t.TempDir() + opts := outputOpts{ + contentDirPrefix: filepath.Join(tmpDir, "contents"), + tempSymlinkPrefix: filepath.Join(tmpDir, "symlink"), + symlink: filepath.Join(tmpDir, "creds"), + } + + err = job.refreshCreds(ctx, opts, time.Now().Format(time.RFC3339)) + if err == nil { + t.Errorf("refreshCreds expected error for 500 internal error, got nil") + } + + if mockServer.reqCount != 1 { + t.Errorf("mockServer.reqCount = %d, want 1 (terminated on 500 internal error)", mockServer.reqCount) + } +} + +func TestRefreshCredsWithGRPCRetry503ThenSuccess(t *testing.T) { + ctx := context.Background() + if err := cfg.Load(nil); err != nil { + t.Fatalf("cfg.Load() failed unexpectedly with error: %v", err) + } + config := cfg.Retrieve() + config.MWLID.Enabled = true + + testPrivateKey := []byte("test-private-key") + testCertChain := []byte("test-cert-chain") + testTrustBundles := []byte("test-trust-bundles") + + mockServer := &mockWorkloadIdentityServer{ + getWorkloadCertificatesResponse: &wipb.GetWorkloadCertificatesResponse{ + CertificateChainPem: testCertChain, PrivateKeyPem: testPrivateKey, + }, + getWorkloadTrustBundlesResponse: &wipb.GetWorkloadTrustBundlesResponse{ + SpiffeTrustBundlesMapJson: testTrustBundles, + }, + getWorkloadCertificatesErrFn: func(reqCount int) error { + if reqCount < 3 { + return status.Error(codes.Unavailable, "temporary 503 unavailable") + } + return nil + }, + } + + addr, stop := startTestGRPCServer(t, mockServer) + defer stop() + + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + t.Fatalf("failed to parse address: %v", err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("failed to parse port: %v", err) + } + + config.MWLID.ServiceIP = host + config.MWLID.ServicePort = port + + job := &RefresherJob{ + retryPolicy: retry.Policy{ + MaxAttempts: 5, + Jitter: time.Millisecond, + }, + } + job.setStatus(ServiceAvailable) + + tmpDir := t.TempDir() + opts := outputOpts{ + contentDirPrefix: filepath.Join(tmpDir, "contents"), + tempSymlinkPrefix: filepath.Join(tmpDir, "symlink"), + symlink: filepath.Join(tmpDir, "creds"), + } + + err = job.refreshCreds(ctx, opts, time.Now().Format(time.RFC3339)) + if err != nil { + t.Errorf("refreshCreds failed unexpectedly: %v", err) + } + + if mockServer.reqCount != 3 { + t.Errorf("mockServer.reqCount = %d, want 3 (succeeded on 3rd attempt after retrying 503)", mockServer.reqCount) + } +} + func TestNewClient(t *testing.T) { if err := cfg.Load(nil); err != nil { t.Fatalf("cfg.Load() failed unexpectedly with error: %v", err)