fix: handle ip prefixes without /32 suffix

qdm12 · Aug 17, 2024 · 3d41cc1 · 3d41cc1
1 parent 7ea1a34
commit 3d41cc1
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 7 deletions.
diff --git a/internal/firewall/list.go b/internal/firewall/list.go
@@ -231,12 +231,12 @@ func parseChainRuleField(fieldIndex int, field string, rule *chainRule) (err err
 	case outputInterfaceIndex:
 		rule.outputInterface = field
 	case sourceIndex:
-		rule.source, err = netip.ParsePrefix(field)
+		rule.source, err = parseIPPrefix(field)
 		if err != nil {
 			return fmt.Errorf("parsing source IP CIDR: %w", err)
 		}
 	case destinationIndex:
-		rule.destination, err = netip.ParsePrefix(field)
+		rule.destination, err = parseIPPrefix(field)
 		if err != nil {
 			return fmt.Errorf("parsing destination IP CIDR: %w", err)
 		}

diff --git a/internal/firewall/list_test.go b/internal/firewall/list_test.go
@@ -56,7 +56,7 @@ num pkts bytes target     prot opt in     out     source               destinati
 num pkts bytes target     prot opt in     out     source               destination
 1   0     0 ACCEPT     17   --  tun0   *       0.0.0.0/0            0.0.0.0/0            udp dpt:55405
 2   0     0 ACCEPT     6    --  tun0   *       0.0.0.0/0            0.0.0.0/0            tcp dpt:55405
-3   0     0 DROP       0    --  tun0   *       0.0.0.0/0            0.0.0.0/0
+3   0     0 DROP       0    --  tun0   *       1.2.3.4              0.0.0.0/0
 `,
 			table: chain{
 				name:    "INPUT",
@@ -96,7 +96,7 @@ num pkts bytes target     prot opt in     out     source               destinati
 						protocol:        "",
 						inputInterface:  "tun0",
 						outputInterface: "*",
-						source:          netip.MustParsePrefix("0.0.0.0/0"),
+						source:          netip.MustParsePrefix("1.2.3.4/32"),
 						destination:     netip.MustParsePrefix("0.0.0.0/0"),
 					},
 				},

diff --git a/internal/firewall/parse.go b/internal/firewall/parse.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"net/netip"
+	"regexp"
 	"slices"
 	"strconv"
 	"strings"
@@ -117,12 +118,12 @@ func parseInstructionFlag(key, value string, instruction *iptablesInstruction) (
 	case "-o", "--out-interface":
 		instruction.outputInterface = value
 	case "-s", "--source":
-		instruction.source, err = netip.ParsePrefix(value)
+		instruction.source, err = parseIPPrefix(value)
 		if err != nil {
 			return fmt.Errorf("parsing source IP CIDR: %w", err)
 		}
 	case "-d", "--destination":
-		instruction.destination, err = netip.ParsePrefix(value)
+		instruction.destination, err = parseIPPrefix(value)
 		if err != nil {
 			return fmt.Errorf("parsing destination IP CIDR: %w", err)
 		}
@@ -151,3 +152,12 @@ func parseInstructionFlag(key, value string, instruction *iptablesInstruction) (
 	}
 	return nil
 }
+
+var regexCidrSuffix = regexp.MustCompile(`/[0-9][0-9]{0,1}$`)
+
+func parseIPPrefix(value string) (prefix netip.Prefix, err error) {
+	if !regexCidrSuffix.MatchString(value) {
+		value += "/32"
+	}
+	return netip.ParsePrefix(value)
+}
diff --git a/internal/firewall/parse_test.go b/internal/firewall/parse_test.go
@@ -39,14 +39,15 @@ func Test_parseIptablesInstruction(t *testing.T) {
 			},
 		},
 		"instruction_A": {
-			s: "-A INPUT -i tun0 -p tcp -m tcp -s 1.2.3.4/32 --dport 10000 -j ACCEPT",
+			s: "-A INPUT -i tun0 -p tcp -m tcp -s 1.2.3.4/32 -d 5.6.7.8 --dport 10000 -j ACCEPT",
 			instruction: iptablesInstruction{
 				table:           "filter",
 				chain:           "INPUT",
 				append:          true,
 				inputInterface:  "tun0",
 				protocol:        "tcp",
 				source:          netip.MustParsePrefix("1.2.3.4/32"),
+				destination:     netip.MustParsePrefix("5.6.7.8/32"),
 				destinationPort: 10000,
 				target:          "ACCEPT",
 			},