diff --git a/sdk/iptables/iptables.go b/sdk/iptables/iptables.go index b2a7a62fb1..7ebd07b8c5 100644 --- a/sdk/iptables/iptables.go +++ b/sdk/iptables/iptables.go @@ -69,18 +69,18 @@ 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) } +// AdditionalRulesFn 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. +type AdditionalRulesFn func(iptablesProvider Provider) + // Provider is an interface for executing iptables rules. type Provider interface { // AddRule adds a rule without executing it. @@ -98,9 +98,15 @@ type Provider interface { // Setup will set up iptables interception and redirection rules // based on the configuration provided in cfg. -// This implementation was inspired by -// https://github.com/openservicemesh/osm/blob/650a1a1dcf081ae90825f3b5dba6f30a0e532725/pkg/injector/iptables.go func Setup(cfg Config) error { + return SetupWithAdditionalRules(cfg, nil) +} + +// SetupWithAdditionalRules will set up iptables interception and redirection rules +// based on the configuration provided in cfg. The additionalRulesFn will be applied +// after the normal set of rules. This implementation was inspired by +// https://github.com/openservicemesh/osm/blob/650a1a1dcf081ae90825f3b5dba6f30a0e532725/pkg/injector/iptables.go +func SetupWithAdditionalRules(cfg Config, additionalRulesFn AdditionalRulesFn) error { if cfg.IptablesProvider == nil { cfg.IptablesProvider = &iptablesExecutor{cfg: cfg} } @@ -193,8 +199,8 @@ 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) + if additionalRulesFn != nil { + additionalRulesFn(cfg.IptablesProvider) } return cfg.IptablesProvider.ApplyRules() diff --git a/sdk/iptables/iptables_test.go b/sdk/iptables/iptables_test.go index b614a45d58..0bd740fd81 100644 --- a/sdk/iptables/iptables_test.go +++ b/sdk/iptables/iptables_test.go @@ -303,15 +303,16 @@ func TestSetup(t *testing.T) { for _, c := range cases { t.Run(c.name, func(t *testing.T) { + var fn AdditionalRulesFn if c.additionalRules != nil { - c.cfg.AddAdditionalRulesFn = func(provider Provider) { + fn = func(provider Provider) { for _, rule := range c.additionalRules { provider.AddRule(rule[0], rule[1:]...) } } } - err := Setup(c.cfg) + err := SetupWithAdditionalRules(c.cfg, fn) require.NoError(t, err) require.Equal(t, c.expectedRules, c.cfg.IptablesProvider.Rules()) })