Address review feedback: nil check, ambiguity detection, test fixes

- Add nil check for c.EC2 in resolveCurrentVPCID to prevent panic
- Detect ambiguous VPC resolution when private IP matches multiple VPCs
- Replace os.Unsetenv with t.Setenv for test isolation
- Replace flaky resolveCurrentVPCID test with nil EC2 client test

https://claude.ai/code/session_019cXyAQaLbKyaHWnZupcNUC

Author: claude

pkg/aws/ec2.go

@@ -1215,6 +1215,10 @@ func (c *EC2Client) tryVPCWideConfiguration(ctx context.Context, hopLimit int32,
 
 // resolveCurrentVPCID determines the VPC ID of the current instance by looking up its private IP
 func (c *EC2Client) resolveCurrentVPCID(ctx context.Context) (string, error) {
+	if c.EC2 == nil {
+		return "", fmt.Errorf("EC2 client is not initialized")
+	}
+
 	privateIP, err := c.getPrivateIPFromNetworkInterface()
 	if err != nil {
 		return "", fmt.Errorf("failed to get private IP for VPC resolution: %v", err)
@@ -1244,15 +1248,30 @@ func (c *EC2Client) resolveCurrentVPCID(ctx context.Context) (string, error) {
 		return "", fmt.Errorf("failed to describe instances for VPC resolution: %v", err)
 	}
 
+	// Collect unique VPC IDs to detect ambiguity from overlapping CIDRs
+	vpcIDs := make(map[string]struct{})
 	for _, reservation := range result.Reservations {
 		for _, instance := range reservation.Instances {
 			if instance.VpcId != nil {
-				c.Logger.V(1).Info("Resolved VPC ID", "vpcID", *instance.VpcId, "privateIP", privateIP)
-				return *instance.VpcId, nil
+				vpcIDs[*instance.VpcId] = struct{}{}
 			}
 		}
 	}
 
+	if len(vpcIDs) == 0 {
+		return "", fmt.Errorf("no VPC ID found for instance with private IP %s", privateIP)
+	}
+
+	if len(vpcIDs) > 1 {
+		return "", fmt.Errorf("ambiguous VPC resolution: private IP %s matched instances in %d different VPCs", privateIP, len(vpcIDs))
+	}
+
+	// Exactly one VPC matched
+	for vpcID := range vpcIDs {
+		c.Logger.V(1).Info("Resolved VPC ID", "vpcID", vpcID, "privateIP", privateIP)
+		return vpcID, nil
+	}
+
 	return "", fmt.Errorf("no VPC ID found for instance with private IP %s", privateIP)
 }
 

pkg/aws/imds_vpc_test.go

@@ -2,7 +2,6 @@ package aws
 
 import (
 	"context"
-	"os"
 	"testing"
 
 	"github.com/go-logr/logr/testr"
@@ -29,8 +28,8 @@ func TestTryVPCWideConfiguration_EmptyVPCID(t *testing.T) {
 // TestTryVPCWideConfiguration_AggressiveDisabled verifies that tryVPCWideConfiguration
 // returns an error when aggressive configuration is disabled.
 func TestTryVPCWideConfiguration_AggressiveDisabled(t *testing.T) {
-	// Ensure aggressive configuration is disabled
-	os.Unsetenv("IMDS_AGGRESSIVE_CONFIGURATION")
+	// t.Setenv automatically restores the original value after the test
+	t.Setenv("IMDS_AGGRESSIVE_CONFIGURATION", "false")
 
 	client := &EC2Client{
 		Logger: testr.New(t),
@@ -47,21 +46,20 @@ func TestTryVPCWideConfiguration_AggressiveDisabled(t *testing.T) {
 	}
 }
 
-// TestResolveCurrentVPCID_NoPrivateIP verifies that resolveCurrentVPCID
-// returns an error when no private IP can be determined.
-func TestResolveCurrentVPCID_NoPrivateIP(t *testing.T) {
-	// Ensure PRIVATE_IP env var is not set so the lookup relies on network interfaces
-	os.Unsetenv("PRIVATE_IP")
-
+// TestResolveCurrentVPCID_NilEC2Client verifies that resolveCurrentVPCID
+// returns a clear error when the EC2 client is not initialized.
+func TestResolveCurrentVPCID_NilEC2Client(t *testing.T) {
 	client := &EC2Client{
 		Logger: testr.New(t),
 	}
 
-	// This will fail because there's no EC2 client and the network interface
-	// lookup will either return a local IP or fail - either way, without a
-	// real EC2 client the DescribeInstances call would fail.
 	_, err := client.resolveCurrentVPCID(context.Background())
 	if err == nil {
-		t.Fatal("expected error when resolving VPC ID without AWS credentials, got nil")
+		t.Fatal("expected error when EC2 client is nil, got nil")
+	}
+
+	expected := "EC2 client is not initialized"
+	if err.Error() != expected {
+		t.Errorf("expected error %q, got %q", expected, err.Error())
 	}
 }