Updated test agent for snat test

test/agent/Dockerfile

@@ -26,8 +26,12 @@ RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 GO111MODULE=on go build \
 RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 GO111MODULE=on go build \
     -a -o metric-server cmd/metric-server/main.go
 
+RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 GO111MODULE=on go build \
+    -a -o snat-utils cmd/snat-utils/main.go
+
 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

@@ -0,0 +1,124 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"net/http"
+	"regexp"
+	"strings"
+	"time"
+
+	"github.com/coreos/go-iptables/iptables"
+)
+
+func main() {
+	var testIPTableRules bool
+	var testExternalDomainConnectivity bool
+	var randomizedSNATValue string
+	var numOfCidrs int
+	var url string
+
+	flag.BoolVar(&testIPTableRules, "testIPTableRules", false, "bool flag when set to true tests validate if IPTable has required rules")
+	flag.StringVar(&randomizedSNATValue, "randomizedSNATValue", "prng", "value for AWS_VPC_K8S_CNI_RANDOMIZESNAT")
+	flag.IntVar(&numOfCidrs, "numOfCidrs", 1, "Number of CIDR blocks in customer VPC")
+	flag.BoolVar(&testExternalDomainConnectivity, "testExternalDomainConnectivity", false, "bool flag when set to true tests if the pod has internet access")
+	flag.StringVar(&url, "url", "https://aws.amazon.com/", "url to check for connectivity")
+
+	flag.Parse()
+
+	if testIPTableRules {
+		err := validateIPTableRules(randomizedSNATValue, numOfCidrs)
+		if err != nil {
+			log.Fatal(err)
+		}
+		log.Printf("Randomized SNAT test passed for AWS_VPC_K8S_CNI_RANDOMIZESNAT: %s\n", randomizedSNATValue)
+	}
+
+	if testExternalDomainConnectivity {
+		err := validateExternalDomainConnectivity(url)
+		if err != nil {
+			log.Fatal(err)
+		}
+		log.Println("External Domain Connectivity test passed")
+	}
+}
+
+func validateExternalDomainConnectivity(url string) error {
+	timeout := time.Duration(120 * time.Second)
+	client := http.Client{
+		Timeout: timeout,
+	}
+	resp, err := client.Get(url)
+	if err != nil {
+		return err
+	}
+
+	if resp.StatusCode != 200 {
+		return fmt.Errorf("%s returned response code: %d", url, resp.StatusCode)
+	}
+	return nil
+}
+
+func validateIPTableRules(randomizedSNATValue string, numOfCidrs int) error {
+	// Check IPTable rules corresponding to AWS_VPC_K8S_CNI_RANDOMIZESNAT
+	expectedString := "random-fully"
+	iptables, err := iptables.New()
+	if err != nil {
+		return err
+	}
+
+	if !iptables.HasRandomFully() || randomizedSNATValue == "hashrandom" {
+		expectedString = "random"
+	}
+
+	containsExpectedString := false
+
+	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
+		}
+		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)
+
+	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", expectedString)
+	} else if randomizedSNATValue != "none" && !containsExpectedString {
+		return fmt.Errorf("failed: did not find expected %s for any of the SNAT rule", expectedString)
+	}
+	return nil
+}

test/agent/go.mod

@@ -3,6 +3,7 @@ module github.com/aws/amazon-vpc-cni-k8s/test/agent
 go 1.14
 
 require (
+	github.com/coreos/go-iptables v0.6.0
 	github.com/vishvananda/netlink v1.1.0
 	golang.org/x/sys v0.0.0-20210426230700-d19ff857e887
 )