1 files changed, 53 insertions, 2 deletions
diff --git a/util/linuxfw/iptables_runner_test.go b/util/linuxfw/iptables_runner_test.go
index 77c753004..b5a13fdba 100644
--- a/util/linuxfw/iptables_runner_test.go
+++ b/util/linuxfw/iptables_runner_test.go
@@ -126,8 +126,6 @@ func TestAddAndDeleteBase(t *testing.T) {
 
 	// Check that the rules were created.
 	tsRulesV4 := []fakeRule{ // table/chain/rule
-		{"filter", "ts-input", []string{"!", "-i", tunname, "-s", tsaddr.ChromeOSVMRange().String(), "-j", "RETURN"}},
-		{"filter", "ts-input", []string{"!", "-i", tunname, "-s", tsaddr.CGNATRange().String(), "-j", "DROP"}},
 		{"filter", "ts-forward", []string{"-o", tunname, "-s", tsaddr.CGNATRange().String(), "-j", "DROP"}},
 	}
 
@@ -504,3 +502,56 @@ func TestAddAndDelConnmarkSaveRule(t *testing.T) {
 		}
 	})
 }
+
+func TestAddAndDelCGNATRules(t *testing.T) {
+	iptr := newFakeIPTablesRunner()
+	tunname := "tun0"
+
+	// We need the chains to exist so we can add rules into them.
+	if err := iptr.AddChains(); err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		mode      CGNATMode
+		wantRules []fakeRule
+	}{
+		{
+			CGNATModeDrop, []fakeRule{
+				{"filter", "ts-input", []string{"!", "-i", tunname, "-s", tsaddr.ChromeOSVMRange().String(), "-j", "RETURN"}},
+				{"filter", "ts-input", []string{"!", "-i", tunname, "-s", tsaddr.CGNATRange().String(), "-j", "DROP"}},
+			},
+		},
+		{
+			CGNATModeReturn, []fakeRule{
+				{"filter", "ts-input", []string{"!", "-i", tunname, "-s", tsaddr.CGNATRange().String(), "-j", "RETURN"}},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		if err := iptr.AddExternalCGNATRules(tt.mode, tunname); err != nil {
+			t.Fatal(err)
+		}
+
+		for _, tr := range tt.wantRules {
+			if exists, err := iptr.ipt4.Exists(tr.table, tr.chain, tr.args...); err != nil {
+				t.Fatalf("mode %q: error checking for rule: %v", tt.mode, err)
+			} else if !exists {
+				t.Errorf("mode %q: rule %s/%s/%s doesn't exist", tt.mode, tr.table, tr.chain, strings.Join(tr.args, " "))
+			}
+		}
+
+		if err := iptr.DelExternalCGNATRules(tt.mode, tunname); err != nil {
+			t.Fatal(err)
+		}
+
+		for _, tr := range tt.wantRules {
+			if exists, err := iptr.ipt4.Exists(tr.table, tr.chain, tr.args...); err != nil {
+				t.Fatalf("mode %q: error checking for rule: %v", tt.mode, err)
+			} else if exists {
+				t.Errorf("mode %q: rule %s/%s/%s not deleted", tt.mode, tr.table, tr.chain, strings.Join(tr.args, " "))
+			}
+		}
+	}
+}