Add iptables rule builder

This patch is used to provide a simple way to build iptables rules.

Signed-off-by: Hongliang Liu <[email protected]>

pkg/agent/util/iptables/builder.go

@@ -0,0 +1,174 @@
+//go:build !windows
+// +build !windows
+
+// Copyright 2023 Antrea Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iptables
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"k8s.io/apimachinery/pkg/util/intstr"
+
+	"antrea.io/antrea/pkg/apis/controlplane/v1beta2"
+)
+
+type iptablesRule struct {
+	chain string
+	specs *strings.Builder
+}
+
+type iptablesRuleBuilder struct {
+	iptablesRule
+}
+
+func NewRuleBuilder(chain string) IPTablesRuleBuilder {
+	builder := &iptablesRuleBuilder{
+		iptablesRule{
+			chain: chain,
+			specs: &strings.Builder{},
+		},
+	}
+	return builder
+}
+
+func (b *iptablesRuleBuilder) MatchIPSetSrc(ipset string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-m set --match-ipset %s src ", ipset)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchIPSetDst(ipset string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-m set --match-ipset %s dst ", ipset)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchTransProtocol(protocol v1beta2.Protocol) IPTablesRuleBuilder {
+	var protoStr string
+	switch protocol {
+	case v1beta2.ProtocolTCP:
+		protoStr = "tcp"
+	case v1beta2.ProtocolUDP:
+		protoStr = "udp"
+	case v1beta2.ProtocolSCTP:
+		protoStr = "sctp"
+	}
+	specStr := fmt.Sprintf("-p %s ", protoStr)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchDstPort(port *intstr.IntOrString, endPort *int32) IPTablesRuleBuilder {
+	if port == nil {
+		return b
+	}
+	var specStr string
+	if endPort != nil {
+		specStr = fmt.Sprintf("--dport %s:%d ", port.String(), *endPort)
+	} else {
+		specStr = fmt.Sprintf("--dport %s ", port.String())
+	}
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchSrcPort(port, endPort *int32) IPTablesRuleBuilder {
+	if port == nil {
+		return b
+	}
+	var matchStr string
+	if endPort != nil {
+		matchStr = fmt.Sprintf("--sport %d:%d ", *port, *endPort)
+	} else {
+		matchStr = fmt.Sprintf("--sport %d ", *port)
+	}
+	b.specs.WriteString(matchStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchICMP(icmpType, icmpCode *int32, ipProtocol Protocol) IPTablesRuleBuilder {
+	parts := []string{"-p"}
+	icmpTypeStr := "icmp"
+	if ipProtocol != ProtocolIPv4 {
+		icmpTypeStr = "icmpv6"
+	}
+	parts = append(parts, icmpTypeStr)
+
+	if icmpType != nil {
+		icmpTypeFlag := "--icmp-type"
+		if ipProtocol != ProtocolIPv4 {
+			icmpTypeFlag = "--icmpv6-type"
+		}
+
+		if icmpCode != nil {
+			parts = append(parts, icmpTypeFlag, fmt.Sprintf("%d/%d", *icmpType, *icmpCode))
+		} else {
+			parts = append(parts, icmpTypeFlag, strconv.Itoa(int(*icmpType)))
+		}
+	}
+	b.specs.WriteString(strings.Join(parts, " "))
+	b.specs.WriteByte(' ')
+
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchInputInterface(interfaceName string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-i %s ", interfaceName)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) MatchOutputInterface(interfaceName string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-o %s ", interfaceName)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) SetTarget(target string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-j %s ", target)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) SetComment(comment string) IPTablesRuleBuilder {
+	specStr := fmt.Sprintf("-m comment --comment %s ", comment)
+	b.specs.WriteString(specStr)
+	return b
+}
+
+func (b *iptablesRuleBuilder) CopyBuilder() IPTablesRuleBuilder {
+	var copiedSpec strings.Builder
+	copiedSpec.Grow(b.specs.Len())
+	copiedSpec.WriteString(b.specs.String())
+	builder := &iptablesRuleBuilder{
+		iptablesRule{
+			chain: b.chain,
+			specs: &copiedSpec,
+		},
+	}
+	return builder
+}
+
+func (b *iptablesRuleBuilder) Done() IPTablesRule {
+	return &b.iptablesRule
+}
+
+func (e *iptablesRule) GetSpec() string {
+	spec := fmt.Sprintf("-A %s %s", e.chain, e.specs)
+	return spec[:len(spec)-1] // Remove the last space.
+}

pkg/agent/util/iptables/builder_test.go

@@ -0,0 +1,123 @@
+//go:build !windows
+// +build !windows
+
+// Copyright 2023 Antrea Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iptables
+
+import (
+	"antrea.io/antrea/pkg/apis/controlplane/v1beta2"
+	"k8s.io/apimachinery/pkg/util/intstr"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	ipsetAlfa  = "alfa"
+	ipsetBravo = "bravo"
+	eth0       = "eth0"
+	eth1       = "eth1"
+	port8080   = &intstr.IntOrString{Type: intstr.Int, IntVal: 8080}
+	port137    = &intstr.IntOrString{Type: intstr.Int, IntVal: 137}
+	port139    = int32(139)
+	port40000  = int32(40000)
+	port50000  = int32(50000)
+	icmpType0  = int32(0)
+	icmpCode0  = int32(0)
+)
+
+func TestBuilders(t *testing.T) {
+	testCases := []struct {
+		name      string
+		chain     string
+		buildFunc func(IPTablesRuleBuilder) IPTablesRuleBuilder
+		expected  string
+	}{
+		{
+			name:  "Accept TCP destination 8080 in FORWARD",
+			chain: ForwardChain,
+			buildFunc: func(builder IPTablesRuleBuilder) IPTablesRuleBuilder {
+				return builder.
+					MatchIPSetSrc(ipsetAlfa).
+					MatchIPSetDst(ipsetBravo).
+					MatchInputInterface(eth0).
+					MatchTransProtocol(v1beta2.ProtocolTCP).
+					MatchDstPort(port8080, nil).
+					SetComment(`"Accept TCP 8080"`).
+					SetTarget(AcceptTarget)
+			},
+			expected: `-A FORWARD -m set --match-ipset alfa src -m set --match-ipset bravo dst -i eth0 -p tcp --dport 8080 -m comment --comment "Accept TCP 8080" -j ACCEPT`,
+		},
+		{
+			name:  "Drop UDP destination 137-139 in INPUT",
+			chain: "INPUT",
+			buildFunc: func(builder IPTablesRuleBuilder) IPTablesRuleBuilder {
+				return builder.
+					MatchIPSetSrc(ipsetAlfa).
+					MatchInputInterface(eth0).
+					MatchTransProtocol(v1beta2.ProtocolUDP).
+					MatchDstPort(port137, &port139).
+					SetComment(`"Drop UDP 137-139"`).
+					SetTarget(DROPTarget)
+			},
+			expected: `-A INPUT -m set --match-ipset alfa src -i eth0 -p udp --dport 137:139 -m comment --comment "Drop UDP 137-139" -j DROP`,
+		},
+		{
+			name:  "Reject SCTP source 40000-50000 in OUTPUT",
+			chain: OutputChain,
+			buildFunc: func(builder IPTablesRuleBuilder) IPTablesRuleBuilder {
+				return builder.
+					MatchOutputInterface(eth1).
+					MatchTransProtocol(v1beta2.ProtocolSCTP).
+					MatchSrcPort(&port40000, &port50000).
+					SetComment(`"Drop SCTP 40000-50000"`).
+					SetTarget(DROPTarget)
+			},
+			expected: `-A OUTPUT -o eth1 -p sctp --sport 40000:50000 -m comment --comment "Drop SCTP 40000-50000" -j DROP`,
+		},
+		{
+			name:  "Accept ICMP IPv4",
+			chain: ForwardChain,
+			buildFunc: func(builder IPTablesRuleBuilder) IPTablesRuleBuilder {
+				return builder.
+					MatchInputInterface(eth0).
+					MatchICMP(&icmpType0, &icmpCode0, ProtocolIPv4).
+					SetTarget(AcceptTarget)
+			},
+			expected: `-A FORWARD -i eth0 -p icmp --icmp-type 0/0 -j ACCEPT`,
+		},
+		{
+			name:  "Accept ICMP IPv6",
+			chain: ForwardChain,
+			buildFunc: func(builder IPTablesRuleBuilder) IPTablesRuleBuilder {
+				return builder.
+					MatchInputInterface(eth0).
+					MatchICMP(&icmpType0, nil, ProtocolIPv6).
+					SetTarget(AcceptTarget)
+			},
+			expected: `-A FORWARD -i eth0 -p icmpv6 --icmpv6-type 0 -j ACCEPT`,
+		},
+	}
+
+	for _, tc := range testCases {
+		builder := NewRuleBuilder(tc.chain)
+		t.Run(tc.name, func(t *testing.T) {
+			copiedBuilder := builder.CopyBuilder()
+			rule := tc.buildFunc(copiedBuilder).Done()
+			assert.Equal(t, tc.expected, rule.GetSpec())
+		})
+	}
+}

pkg/agent/util/iptables/interfaces.go

@@ -0,0 +1,43 @@
+//go:build !windows
+// +build !windows
+
+// Copyright 2023 Antrea Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iptables
+
+import (
+	"k8s.io/apimachinery/pkg/util/intstr"
+
+	"antrea.io/antrea/pkg/apis/controlplane/v1beta2"
+)
+
+type IPTablesRuleBuilder interface {
+	MatchIPSetSrc(ipset string) IPTablesRuleBuilder
+	MatchIPSetDst(ipset string) IPTablesRuleBuilder
+	MatchTransProtocol(protocol v1beta2.Protocol) IPTablesRuleBuilder
+	MatchDstPort(port *intstr.IntOrString, endPort *int32) IPTablesRuleBuilder
+	MatchSrcPort(port, endPort *int32) IPTablesRuleBuilder
+	MatchICMP(icmpType, icmpCode *int32, ipProtocol Protocol) IPTablesRuleBuilder
+	MatchInputInterface(interfaceName string) IPTablesRuleBuilder
+	MatchOutputInterface(interfaceName string) IPTablesRuleBuilder
+	SetTarget(target string) IPTablesRuleBuilder
+	SetComment(comment string) IPTablesRuleBuilder
+	CopyBuilder() IPTablesRuleBuilder
+	Done() IPTablesRule
+}
+
+type IPTablesRule interface {
+	GetSpec() string
+}