From ff99bfd4a674b18a0ad3bcdec15cf112bbec68d7 Mon Sep 17 00:00:00 2001
From: pufferffish <github@bandersnatch.anonaddy.com>
Date: Mon, 22 Jul 2024 15:10:12 +0100
Subject: [PATCH] fix config parsing

---
 config.go      | 16 ++++++----
 config_test.go | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++
 wireguard.go   |  6 ++--
 3 files changed, 100 insertions(+), 9 deletions(-)
 create mode 100644 config_test.go

diff --git a/config.go b/config.go
index 76593cf..b70bfaf 100644
--- a/config.go
+++ b/config.go
@@ -149,13 +149,17 @@ func parseCIDRNetIP(section *ini.Section, keyName string) ([]netip.Addr, error)
 	keys := key.StringsWithShadows(",")
 	var ips = make([]netip.Addr, 0, len(keys))
 	for _, str := range keys {
-		prefix, err := netip.ParsePrefix(str)
-		if err != nil {
-			return nil, err
-		}
+		if addr, err := netip.ParseAddr(str); err == nil {
+			ips = append(ips, addr)
+		} else {
+			prefix, err := netip.ParsePrefix(str)
+			if err != nil {
+				return nil, err
+			}
 
-		addr := prefix.Addr()
-		ips = append(ips, addr)
+			addr := prefix.Addr()
+			ips = append(ips, addr)
+		}
 	}
 	return ips, nil
 }
diff --git a/config_test.go b/config_test.go
new file mode 100644
index 0000000..948fbf8
--- /dev/null
+++ b/config_test.go
@@ -0,0 +1,87 @@
+package wireproxy
+
+import (
+	"github.com/go-ini/ini"
+	"testing"
+)
+
+func loadIniConfig(config string) (*ini.File, error) {
+	iniOpt := ini.LoadOptions{
+		Insensitive:            true,
+		AllowShadows:           true,
+		AllowNonUniqueSections: true,
+	}
+
+	return ini.LoadSources(iniOpt, []byte(config))
+}
+
+func TestWireguardConfWithoutSubnet(t *testing.T) {
+	const config = `
+[Interface]
+PrivateKey = LAr1aNSNF9d0MjwUgAVC4020T0N/E5NUtqVv5EnsSz0=
+Address = 10.5.0.2
+DNS = 1.1.1.1
+
+[Peer]
+PublicKey = e8LKAc+f9xEzq9Ar7+MfKRrs+gZ/4yzvpRJLRJ/VJ1w=
+AllowedIPs = 0.0.0.0/0, ::/0
+Endpoint = 94.140.11.15:51820
+PersistentKeepalive = 25`
+	var cfg DeviceConfig
+	iniData, err := loadIniConfig(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = ParseInterface(iniData, &cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestWireguardConfWithSubnet(t *testing.T) {
+	const config = `
+[Interface]
+PrivateKey = LAr1aNSNF9d0MjwUgAVC4020T0N/E5NUtqVv5EnsSz0=
+Address = 10.5.0.2/23
+DNS = 1.1.1.1
+
+[Peer]
+PublicKey = e8LKAc+f9xEzq9Ar7+MfKRrs+gZ/4yzvpRJLRJ/VJ1w=
+AllowedIPs = 0.0.0.0/0, ::/0
+Endpoint = 94.140.11.15:51820
+PersistentKeepalive = 25`
+	var cfg DeviceConfig
+	iniData, err := loadIniConfig(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = ParseInterface(iniData, &cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestWireguardConfWithManyAddress(t *testing.T) {
+	const config = `
+[Interface]
+PrivateKey = mBsVDahr1XIu9PPd17UmsDdB6E53nvmS47NbNqQCiFM=
+Address = 100.96.0.190,2606:B300:FFFF:fe8a:2ac6:c7e8:b021:6f5f/128
+DNS = 198.18.0.1,198.18.0.2
+
+[Peer]
+PublicKey = SHnh4C2aDXhp1gjIqceGhJrhOLSeNYcqWLKcYnzj00U=
+AllowedIPs = 0.0.0.0/0,::/0
+Endpoint = 192.200.144.22:51820`
+	var cfg DeviceConfig
+	iniData, err := loadIniConfig(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = ParseInterface(iniData, &cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/wireguard.go b/wireguard.go
index 31057ed..8b2d0f8 100644
--- a/wireguard.go
+++ b/wireguard.go
@@ -20,8 +20,8 @@ type DeviceSetting struct {
 	mtu        int
 }
 
-// serialize the config into an IPC request and DeviceSetting
-func createIPCRequest(conf *DeviceConfig) (*DeviceSetting, error) {
+// CreateIPCRequest serialize the config into an IPC request and DeviceSetting
+func CreateIPCRequest(conf *DeviceConfig) (*DeviceSetting, error) {
 	var request bytes.Buffer
 
 	request.WriteString(fmt.Sprintf("private_key=%s\n", conf.SecretKey))
@@ -60,7 +60,7 @@ func createIPCRequest(conf *DeviceConfig) (*DeviceSetting, error) {
 
 // StartWireguard creates a tun interface on netstack given a configuration
 func StartWireguard(conf *DeviceConfig, logLevel int) (*VirtualTun, error) {
-	setting, err := createIPCRequest(conf)
+	setting, err := CreateIPCRequest(conf)
 	if err != nil {
 		return nil, err
 	}