Add additional rules to iptables SDK (#20232)

* Add additional rules to iptables SDK

* Add comments

sdk/iptables/iptables.go

@@ -69,6 +69,16 @@ type Config struct {
 
 	// IptablesProvider is the Provider that will apply iptables rules.
 	IptablesProvider Provider
+
+	// AddAdditionalRulesFn can be implemented by the caller to
+	// add environment specific rules (like ECS) that needs to
+	// be executed for traffic redirection to work properly.
+	//
+	// This gets called by the Setup function after all the
+	// first class iptable rules are added. The implemented
+	// function should only call the `AddRule` and optionally
+	// the `Rules` method of the provider.
+	AddAdditionalRulesFn func(iptablesProvider Provider)
 }
 
 // Provider is an interface for executing iptables rules.
@@ -182,6 +192,11 @@ func Setup(cfg Config) error {
 		}
 	}
 
+	// Call function to add any additional rules passed on by the caller
+	if cfg.AddAdditionalRulesFn != nil {
+		cfg.AddAdditionalRulesFn(cfg.IptablesProvider)
+	}
+
 	return cfg.IptablesProvider.ApplyRules()
 }
 

sdk/iptables/iptables_test.go

@@ -12,9 +12,10 @@ import (
 
 func TestSetup(t *testing.T) {
 	cases := []struct {
-		name          string
+		name            string
-		cfg           Config
+		cfg             Config
-		expectedRules []string
+		additionalRules [][]string
+		expectedRules   []string
 	}{
 		{
 			"no proxy outbound port provided",
@@ -23,6 +24,7 @@ func TestSetup(t *testing.T) {
 				ProxyInboundPort: 20000,
 				IptablesProvider: &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -47,6 +49,7 @@ func TestSetup(t *testing.T) {
 				ConsulDNSIP:      "10.0.34.16",
 				IptablesProvider: &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -75,6 +78,7 @@ func TestSetup(t *testing.T) {
 				ConsulDNSPort:    8600,
 				IptablesProvider: &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -104,6 +108,7 @@ func TestSetup(t *testing.T) {
 				ConsulDNSPort:    8600,
 				IptablesProvider: &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -132,6 +137,7 @@ func TestSetup(t *testing.T) {
 				ProxyOutboundPort: 21000,
 				IptablesProvider:  &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -157,6 +163,7 @@ func TestSetup(t *testing.T) {
 				ExcludeInboundPorts: []string{"22000", "22500"},
 				IptablesProvider:    &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -184,6 +191,7 @@ func TestSetup(t *testing.T) {
 				ExcludeOutboundPorts: []string{"22000", "22500"},
 				IptablesProvider:     &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -211,6 +219,7 @@ func TestSetup(t *testing.T) {
 				ExcludeOutboundCIDRs: []string{"1.1.1.1", "2.2.2.2/24"},
 				IptablesProvider:     &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -238,6 +247,7 @@ func TestSetup(t *testing.T) {
 				ExcludeUIDs:       []string{"456", "789"},
 				IptablesProvider:  &fakeIptablesProvider{},
 			},
+			nil,
 			[]string{
 				"iptables -t nat -N CONSUL_PROXY_INBOUND",
 				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
@@ -256,10 +266,51 @@ func TestSetup(t *testing.T) {
 				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
 			},
 		},
+		{
+			"additional rules are passed",
+			Config{
+				ProxyUserID:       "123",
+				ProxyInboundPort:  20000,
+				ProxyOutboundPort: 21000,
+				ExcludeUIDs:       []string{"456", "789"},
+				IptablesProvider:  &fakeIptablesProvider{},
+			},
+			[][]string{
+				{"iptables", "-t", "nat", "--policy", "POSTROUTING", "ACCEPT"},
+				{"iptables", "-t", "nat", "--policy", "PREROUTING", "ACCEPT"},
+			},
+			[]string{
+				"iptables -t nat -N CONSUL_PROXY_INBOUND",
+				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
+				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
+				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
+				"iptables -t nat -N CONSUL_DNS_REDIRECT",
+				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
+				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
+				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
+				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
+				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
+				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 456 -j RETURN",
+				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 789 -j RETURN",
+				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
+				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
+				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
+				"iptables -t nat --policy POSTROUTING ACCEPT",
+				"iptables -t nat --policy PREROUTING ACCEPT",
+			},
+		},
 	}
 
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
+			if c.additionalRules != nil {
+				c.cfg.AddAdditionalRulesFn = func(provider Provider) {
+					for _, rule := range c.additionalRules {
+						provider.AddRule(rule[0], rule[1:]...)
+					}
+				}
+			}
+
 			err := Setup(c.cfg)
 			require.NoError(t, err)
 			require.Equal(t, c.expectedRules, c.cfg.IptablesProvider.Rules())