Fix parser error causing failure of all tests

NHAS · Oct 31, 2024 · d214440 · d214440
1 parent d1888c7
commit d214440
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 15 deletions.
diff --git a/internal/config/config.go b/internal/config/config.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"net/netip"
 	"os"
 	"strconv"
 	"strings"
@@ -365,9 +366,9 @@ func Load(path string) error {
 }
 
 func parseAddress(address string) ([]string, error) {
-	ip := net.ParseIP(address)
-	if ip == nil {
 
+	addr, err := netip.ParseAddr(address)
+	if err != nil {
 		_, cidr, err := net.ParseCIDR(address)
 		if err != nil {
 
@@ -406,9 +407,9 @@ func parseAddress(address string) ([]string, error) {
 	}
 
 	mask := "/32"
-	if ip.To16() == nil {
+	if addr.Is6() {
 		mask = "/128"
 	}
 
-	return []string{ip.String() + mask}, nil
+	return []string{addr.String() + mask}, nil
 }
diff --git a/internal/router/wireguard.go b/internal/router/wireguard.go
@@ -337,10 +337,13 @@ func (f *Firewall) setupUsers(users []data.UserModel) error {
 }
 
 func (f *Firewall) hostIPWithMask(s string) string {
-	ip := net.ParseIP(s)
+	addr, err := netip.ParseAddr(s)
+	if err != nil {
+		return s + "/32"
+	}
 
 	mask := "/32"
-	if ip.To16() == nil {
+	if addr.Is6() {
 		mask = "/128"
 	}
 

diff --git a/internal/routetypes/key.go b/internal/routetypes/key.go
@@ -29,6 +29,10 @@ func (l *Key) AsIPv6() net.IP {
 	return net.IP(l.IP).To16()
 }
 
+func (l *Key) AsIP() net.IP {
+	return net.IP(l.IP)
+}
+
 func (l Key) String() string {
 	return fmt.Sprintf("%s/%d", net.IP(l.IP).String(), l.Prefixlen)
 }

diff --git a/internal/routetypes/parser.go b/internal/routetypes/parser.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/netip"
 	"sort"
 	"strconv"
 	"strings"
@@ -359,8 +360,8 @@ var (
 
 func parseAddress(address string) (resultAddresses []net.IPNet, err error) {
 
-	ip := net.ParseIP(address)
-	if ip == nil {
+	addr, err := netip.ParseAddr(address)
+	if err != nil {
 
 		_, cidr, err := net.ParseCIDR(address)
 		if err != nil {
@@ -412,13 +413,13 @@ func parseAddress(address string) (resultAddresses []net.IPNet, err error) {
 	}
 
 	var resultIP net.IPNet
-
-	if ip.To4() == nil {
-		resultIP.IP = ip.To4()
+	resultIP.IP = addr.AsSlice()
+	if addr.Is4() {
 		resultIP.Mask = net.CIDRMask(32, 32)
-	} else if ip.To16() == nil {
-		resultIP.IP = ip.To16()
+	} else if addr.Is6() {
 		resultIP.Mask = net.CIDRMask(128, 128)
+	} else {
+		return nil, errors.New("invalid ip")
 	}
 
 	return []net.IPNet{

diff --git a/internal/routetypes/parser_test.go b/internal/routetypes/parser_test.go
@@ -8,8 +8,8 @@ import (
 
 func checkKey(reality Key, expectedKey Key) error {
 
-	if !net.IP.Equal(reality.AsIPv4(), expectedKey.AsIPv4()) {
-		return fmt.Errorf("key had incorrect ip: expected: %s got: %s", expectedKey.IP, reality.IP)
+	if !net.IP.Equal(reality.AsIP(), expectedKey.AsIP()) {
+		return fmt.Errorf("key had incorrect ip: expected: %s got: %s", expectedKey.AsIP(), reality.AsIP())
 	}
 
 	if reality.Prefixlen != expectedKey.Prefixlen {