Updated snat rule test logic

Install iptables in the test agent image

test/agent/Dockerfile

@@ -31,6 +31,7 @@ RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 GO111MODULE=on go build \
 
 FROM public.ecr.aws/amazonlinux/amazonlinux:2
 RUN yum update -y && \
+    yum install -y iptables && \
     yum clean all
 
 WORKDIR /

test/agent/cmd/snat-utils/main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"regexp"
 	"strings"
 	"time"
 
@@ -72,33 +73,52 @@ func validateIPTableRules(randomizedSNATValue string, numOfCidrs int) error {
 	}
 
 	containsExpectedString := false
-	rule := ""
 
-	for i := 0; i <= numOfCidrs; i++ {
-		curr := fmt.Sprintf("AWS-SNAT-CHAIN-%d", i)
-		fmt.Printf("Checking: %s\n", curr)
-		chains, err := iptables.List("nat", curr)
+	currChain := "AWS-SNAT-CHAIN-0"
+	lastChain := fmt.Sprintf("AWS-SNAT-CHAIN-%d", numOfCidrs)
+	i := 0
+	for i < numOfCidrs {
+		rules, err := iptables.List("nat", currChain)
 		if err != nil {
 			return err
 		}
-
-		for _, chain := range chains {
-			if strings.Contains(chain, expectedString) {
-				rule = chain
-				containsExpectedString = true
+		i = i + 1
+		nextChain := fmt.Sprintf("AWS-SNAT-CHAIN-%d", i)
+		foundNextChain := false
+		for _, rule := range rules {
+			target := fmt.Sprintf("-j %s", nextChain)
+			if strings.Contains(rule, target) {
+				currChain = nextChain
+				foundNextChain = true
 				break
 			}
 		}
+		if foundNextChain == false {
+			return fmt.Errorf("failed: AWS-SNAT chain broken for %s", currChain)
+		}
+	}
+
+	// Fetch rules from lastChain
+	rules, err := iptables.List("nat", lastChain)
+	if err != nil {
+		return err
+	}
+
+	// Check for rule with following pattern
+	match := fmt.Sprintf(".*-j SNAT.*%s", expectedString)
+	r, _ := regexp.Compile(match)
 
-		if containsExpectedString {
+	for _, rule := range rules {
+		if r.Match([]byte(rule)) {
+			containsExpectedString = true
 			break
 		}
 	}
 
 	if randomizedSNATValue == "none" && containsExpectedString {
-		return fmt.Errorf("failed: found unexpected %s for SNAT rule: %s", expectedString, rule)
+		return fmt.Errorf("failed: found unexpected %s for SNAT rule", expectedString)
 	} else if randomizedSNATValue != "none" && !containsExpectedString {
-		return fmt.Errorf("failed: did not find expected %s for any of the SNAT rules", expectedString)
+		return fmt.Errorf("failed: did not find expected %s for any of the SNAT rule", expectedString)
 	}
 	return nil
 }