adelorenzo/kilo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make usage of RuleSet prettier

Browse Source
This commit is contained in:
Alex Stockinger
2022-07-26 13:45:57 +00:00
parent 378dafffe8
commit 46cdd6c60c
6 changed files with 59 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
pkg/iptables/iptables.go
View File

@@ -46,6 +46,11 @@ func ipv6Disabled() (bool, error) {
// Protocol represents an IP protocol.
type Protocol byte
type RuleSet struct {
appendRules []Rule // Rules to append to the chain - order matters.
prependRules []Rule // Rules to prepend to the chain - order does not matter.
}
const (
// ProtocolIPv4 represents the IPv4 protocol.
ProtocolIPv4 Protocol = iota
@@ -53,6 +58,21 @@ const (
ProtocolIPv6
)
func (rs *RuleSet) AddToAppend(rule Rule) {
rs.appendRules = append(rs.appendRules, rule)
}
func (rs *RuleSet) AddToPrepend(rule Rule) {
rs.prependRules = append(rs.prependRules, rule)
}
func (rs *RuleSet) AppendRuleSet(other RuleSet) RuleSet {
return RuleSet{
appendRules: append(rs.appendRules, other.appendRules...),
prependRules: append(rs.prependRules, other.prependRules...),
}
}
// GetProtocol will return a protocol from the length of an IP address.
func GetProtocol(ip net.IP) Protocol {
if len(ip) == net.IPv4len || ip.To4() != nil {
@@ -423,10 +443,10 @@ func (c *Controller) deleteFromIndex(i int, rules *[]Rule) error {
func (c *Controller) Set(rules RuleSet) error {
c.Lock()
defer c.Unlock()
if err := c.setAppendRules(rules.AppendRules); err != nil {
if err := c.setAppendRules(rules.appendRules); err != nil {
return err
}
return c.setPrependRules(rules.PrependRules)
return c.setPrependRules(rules.prependRules)
}
func (c *Controller) setAppendRules(appendRules []Rule) error {
@@ -520,8 +540,3 @@ func nonBlockingSend(errors chan<- error, err error) {
default:
}
}
type RuleSet struct {
AppendRules []Rule // Rules to append to the chain - order matters.
PrependRules []Rule // Rules to prepend to the chain - order does not matter.
}

20
pkg/iptables/iptables_test.go
View File

@@ -43,7 +43,7 @@ func TestSet(t *testing.T) {
{
name: "single",
sets: []RuleSet{
{AppendRules: []Rule{appendRules[0]}},
{appendRules: []Rule{appendRules[0]}},
},
appendOut: []Rule{appendRules[0]},
storageOut: []Rule{appendRules[0]},
@@ -51,7 +51,7 @@ func TestSet(t *testing.T) {
{
name: "two rules",
sets: []RuleSet{
{AppendRules: []Rule{appendRules[0], appendRules[1]}},
{appendRules: []Rule{appendRules[0], appendRules[1]}},
},
appendOut: []Rule{appendRules[0], appendRules[1]},
storageOut: []Rule{appendRules[0], appendRules[1]},
@@ -59,8 +59,8 @@ func TestSet(t *testing.T) {
{
name: "multiple",
sets: []RuleSet{
{AppendRules: []Rule{appendRules[0], appendRules[1]}},
{AppendRules: []Rule{appendRules[1]}},
{appendRules: []Rule{appendRules[0], appendRules[1]}},
{appendRules: []Rule{appendRules[1]}},
},
appendOut: []Rule{appendRules[1]},
storageOut: []Rule{appendRules[1]},
@@ -68,7 +68,7 @@ func TestSet(t *testing.T) {
{
name: "re-add",
sets: []RuleSet{
{AppendRules: []Rule{appendRules[0], appendRules[1]}},
{appendRules: []Rule{appendRules[0], appendRules[1]}},
},
appendOut: []Rule{appendRules[0], appendRules[1]},
storageOut: []Rule{appendRules[0], appendRules[1]},
@@ -84,7 +84,7 @@ func TestSet(t *testing.T) {
{
name: "order",
sets: []RuleSet{
{AppendRules: []Rule{appendRules[0], appendRules[1]}},
{appendRules: []Rule{appendRules[0], appendRules[1]}},
},
appendOut: []Rule{appendRules[0], appendRules[1]},
storageOut: []Rule{appendRules[0], appendRules[1]},
@@ -98,8 +98,8 @@ func TestSet(t *testing.T) {
name: "append and prepend",
sets: []RuleSet{
{
PrependRules: []Rule{prependRules[0], prependRules[1]},
AppendRules: []Rule{appendRules[0], appendRules[1]},
prependRules: []Rule{prependRules[0], prependRules[1]},
appendRules: []Rule{appendRules[0], appendRules[1]},
},
},
appendOut: []Rule{appendRules[0], appendRules[1]},
@@ -184,12 +184,12 @@ func TestCleanUp(t *testing.T) {
if err != nil {
t.Fatalf("test case %q: got unexpected error instantiating controller: %v", tc.name, err)
}
ruleSet := RuleSet{AppendRules: tc.appendRules, PrependRules: tc.prependRules}
ruleSet := RuleSet{appendRules: tc.appendRules, prependRules: tc.prependRules}
if err := controller.Set(ruleSet); err != nil {
t.Fatalf("test case %q: Set should not fail: %v", tc.name, err)
}
if len(client.storage) != len(tc.appendRules)+len(tc.prependRules) {
t.Errorf("test case %q: expected %d rules in storage, got %d rules", tc.name, len(ruleSet.AppendRules)+len(ruleSet.PrependRules), len(client.storage))
t.Errorf("test case %q: expected %d rules in storage, got %d rules", tc.name, len(ruleSet.appendRules)+len(ruleSet.prependRules), len(client.storage))
}
if err := controller.CleanUp(); err != nil {
t.Errorf("test case %q: got unexpected error: %v", tc.name, err)

2
pkg/iptables/rulecache_test.go
View File

@@ -101,7 +101,7 @@ func TestRuleCache(t *testing.T) {
client := &fakeClient{}
controller.v4 = client
controller.v6 = client
ruleSet := RuleSet{AppendRules: tc.rules}
ruleSet := RuleSet{appendRules: tc.rules}
if err := controller.Set(ruleSet); err != nil {
t.Fatalf("test case %q: Set should not fail: %v", tc.name, err)
}