Add validation for IPv6 addresses passed to Fault Injection APIs

ecs-agent/tmds/handlers/fault/v1/types/types.go

@@ -19,6 +19,8 @@ import (
 	"net"
 	"strconv"
 
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
 	"github.com/aws/aws-sdk-go/aws"
 )
 
@@ -201,19 +203,29 @@ func NewNetworkFaultInjectionErrorResponse(err string) NetworkFaultInjectionResp
 
 func validateNetworkFaultRequestSources(sources []*string, sourcesType string) error {
 	for _, element := range sources {
-		elementStr := aws.StringValue(element)
-		validIp := true
-		if net.ParseIP(elementStr) == nil {
-			validIp = false
-		}
-		validIpCIDRBlock := true
-		if _, _, err := net.ParseCIDR(elementStr); err != nil {
-			validIpCIDRBlock = false
-		}
-
-		if !validIpCIDRBlock && !validIp {
-			return fmt.Errorf(invalidValueError, elementStr, sourcesType)
+		if err := validateNetworkFaultRequestSource(aws.StringValue(element), sourcesType); err != nil {
+			return err
 		}
 	}
 	return nil
 }
+
+func validateNetworkFaultRequestSource(source string, sourceType string) error {
+	ip := net.ParseIP(source)
+	if ip != nil && ip.To4() != nil {
+		return nil // IPv4 successful
+	}
+
+	_, ipnet, err := net.ParseCIDR(source)
+	if err == nil && ipnet.IP.To4() != nil {
+		return nil // IPv4 CIDR successful
+	}
+	if err != nil {
+		logger.Info("Failed to parse fault source as IPv4 CIDR block", logger.Fields{
+			"source":    source,
+			field.Error: err,
+		})
+	}
+
+	return fmt.Errorf(invalidValueError, source, sourceType)
+}

ecs-agent/tmds/handlers/fault/v1/types/types_test.go

@@ -14,6 +14,7 @@
 package types
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -41,3 +42,29 @@ func TestNetworkBlackholePortAddSourceToFilterIfNotAlready(t *testing.T) {
 		require.Equal(t, aws.StringValueSlice(req.SourcesToFilter), []string{"8.8.8.8", "1.2.3.4"})
 	})
 }
+
+// Tests for validateNetworkFaultRequestSource function that parses IPv4 and IPv4 CIDR blocks.
+func TestValidateNetworkFaultRequestSources(t *testing.T) {
+	tcs := []struct {
+		Name          string
+		Input         string
+		ShouldSucceed bool
+	}{
+		{"IPv4", "1.2.3.4", true},
+		{"IPv4 CIDR", "1.2.3.4/10", true},
+		{"IPv6", "2001:db8::68", false},
+		{"IPv6 CIDR", "::1/128", false},
+		{"invalid input", "invalid", false},
+		{"empty input", "", false},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.Name, func(t *testing.T) {
+			err := validateNetworkFaultRequestSource(tc.Input, "input")
+			if tc.ShouldSucceed {
+				require.NoError(t, err)
+			} else {
+				require.EqualError(t, err, fmt.Sprintf("invalid value %s for parameter input", tc.Input))
+			}
+		})
+	}
+}