Read allow keys from a file, reload on SIGHUP

README.md

@@ -137,6 +137,14 @@ wpex --allow AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= \
   --allow BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB=
 ```
 
+Alternatively, you can point wpex to your wireguard configuration file, or any file with valid `[Peer]` sections.
+
+```bash
+wpex --allow-from-conf /etc/wireguard/wg0.conf
+```
+
+Sending a `SIGHUP` signal to the `wpex` process will trigger a reload of the allowed public keys list from that file.
+
 ## How `wpex` Works
 
 Within each WireGuard session, every peer in the session selects a random 32-bit

go.mod

@@ -7,4 +7,7 @@ require (
 	golang.org/x/time v0.3.0
 )
 
-require golang.org/x/sys v0.12.0 // indirect
+require (
+	golang.org/x/sys v0.12.0 // indirect
+	gopkg.in/ini.v1 v1.67.3 // indirect
+)

go.sum

@@ -1,6 +1,22 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 golang.org/x/crypto v0.13.0 h1:mvySKfSWJ+UKUii46M40LOvyWfN0s2U+46/jDd0e6Ck=
 golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
 golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
 golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/ini.v1 v1.67.3 h1:iM9Lhz5MRSGhHVGGwCuzG9KO8PoirCXj/m/qTmOJJQw=
+gopkg.in/ini.v1 v1.67.3/go.mod h1:x/cyOwCgZqOkJoDIJ3c1KNHMo10+nLGAhh+kn3Zizss=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/analyzer/analyzer.go

@@ -3,11 +3,12 @@ package analyzer
 import (
 	"encoding/binary"
 	"fmt"
-	"github.com/weiiwang01/wpex/internal/exchange"
 	"log"
 	"log/slog"
 	"net"
 	"time"
+
+	"github.com/weiiwang01/wpex/internal/exchange"
 )
 
 const (
@@ -179,12 +180,17 @@ func MakeWireguardAnalyzer(pubkeys [][]byte) WireguardAnalyzer {
 	if err != nil {
 		log.Fatal(fmt.Errorf("failed to generate cookie secret: %w", err))
 	}
-	return WireguardAnalyzer{
+	analyzer := WireguardAnalyzer{
 		table: exchange.MakeExchangeTable(),
 		checker: macChecker{
-			pubkeys: pubkeys,
-			secret:  [32]byte(secret),
-			start:   time.Now(),
+			secret: [32]byte(secret),
+			start:  time.Now(),
 		},
 	}
+	analyzer.SetPubkeys(pubkeys)
+	return analyzer
+}
+
+func (t *WireguardAnalyzer) SetPubkeys(pubkeys [][]byte) {
+	t.checker.SetPubkeys(pubkeys)
 }

internal/analyzer/analyzer_test.go

@@ -45,6 +45,19 @@ func TestWireguardAnalyzer_VerifyMac1(t *testing.T) {
 	}
 }
 
+func TestWireguardAnalyzer_SetPubkeys(t *testing.T) {
+	analyzer := MakeWireguardAnalyzer([][]byte{fakePubkey})
+	addr1, _ := net.ResolveUDPAddr("udp", "127.0.0.1:51820")
+	if _, data := analyzer.Analyse(handshakeInitiationSession1AtoB, *addr1); data != nil {
+		t.Fatal("expected fake pubkey to reject packet")
+	}
+
+	analyzer.SetPubkeys([][]byte{pubkeyA, pubkeyB})
+	if _, data := analyzer.Analyse(handshakeInitiationSession1AtoB, *addr1); data == nil {
+		t.Fatal("expected reloaded pubkeys to accept packet")
+	}
+}
+
 func TestWireguardAnalyzer_Handshake(t *testing.T) {
 	analyzer := MakeWireguardAnalyzer([][]byte{})
 

internal/analyzer/checker.go

@@ -4,13 +4,37 @@ import (
 	"crypto/hmac"
 	"encoding/binary"
 	"net"
+	"sync/atomic"
 	"time"
 )
 
 type macChecker struct {
 	secret  [32]byte
 	start   time.Time
-	pubkeys [][]byte
+	pubkeys atomic.Value
+}
+
+func clonePubkeys(pubkeys [][]byte) [][]byte {
+	if pubkeys == nil {
+		return nil
+	}
+	cloned := make([][]byte, len(pubkeys))
+	for i, key := range pubkeys {
+		cloned[i] = append([]byte(nil), key...)
+	}
+	return cloned
+}
+
+func (c *macChecker) SetPubkeys(pubkeys [][]byte) {
+	c.pubkeys.Store(clonePubkeys(pubkeys))
+}
+
+func (c *macChecker) currentPubkeys() [][]byte {
+	snapshot, ok := c.pubkeys.Load().([][]byte)
+	if !ok {
+		return nil
+	}
+	return snapshot
 }
 
 func (c *macChecker) cookie(addr net.UDPAddr) [16]byte {
@@ -29,13 +53,14 @@ func (c *macChecker) MatchPubkey(message []byte) []byte {
 	if len(message) < 2*macSize {
 		return nil
 	}
-	if len(c.pubkeys) == 0 {
+	pubkeys := c.currentPubkeys()
+	if len(pubkeys) == 0 {
 		return nil
 	}
 	l := len(message)
 	mac1 := message[l-2*macSize : l-macSize]
 	d := message[:l-2*macSize]
-	for _, pubkey := range c.pubkeys {
+	for _, pubkey := range pubkeys {
 		key := hash(nil, []byte("mac1----"), pubkey)
 		m := mac32(nil, key, d)
 		if hmac.Equal(mac1, m[:]) {
@@ -62,5 +87,5 @@ func (c *macChecker) CreateReply(pubkey []byte, addr net.UDPAddr, msg []byte) ([
 }
 
 func (c *macChecker) RequireCheck() bool {
-	return len(c.pubkeys) > 0
+	return len(c.currentPubkeys()) > 0
 }

internal/relay/relay.go

@@ -3,13 +3,16 @@ package relay
 import (
 	"context"
 	"fmt"
-	"github.com/weiiwang01/wpex/internal/analyzer"
-	"golang.org/x/time/rate"
 	"log"
 	"log/slog"
 	"net"
+	"os"
+	"os/signal"
 	"runtime"
 	"syscall"
+
+	"github.com/weiiwang01/wpex/internal/analyzer"
+	"golang.org/x/time/rate"
 )
 
 type udpPacket struct {
@@ -51,7 +54,7 @@ func (r *Relay) relay(conn *net.UDPConn) {
 }
 
 // Start starts the wireguard packet relay server.
-func Start(address string, publicKeys [][]byte, broadcastLimit *rate.Limiter) {
+func Start(address string, getPublicKeys func() ([][]byte, error), broadcastLimit *rate.Limiter) {
 	slog.Info("server listening", "addr", address)
 	var lc = net.ListenConfig{
 		Control: func(network, address string, c syscall.RawConn) error {
@@ -62,11 +65,32 @@ func Start(address string, publicKeys [][]byte, broadcastLimit *rate.Limiter) {
 			return opErr
 		},
 	}
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGHUP)
+
+	initialKeys, err := getPublicKeys()
+	if err != nil {
+		log.Fatalf("failed to get allowed keys: %s", err)
+	}
+
 	relay := Relay{
 		send:     make(chan udpPacket),
-		analyzer: analyzer.MakeWireguardAnalyzer(publicKeys),
+		analyzer: analyzer.MakeWireguardAnalyzer(initialKeys),
 		limit:    broadcastLimit,
 	}
+	go func() {
+		for sig := range sigChan {
+			if sig == syscall.SIGHUP {
+				slog.Info("received SIGHUP, reloading public keys")
+				newKeys, err := getPublicKeys()
+				if err != nil {
+					slog.Error("failed to get allowed keys", "error", err)
+					continue
+				}
+				relay.analyzer.SetPubkeys(newKeys)
+			}
+		}
+	}()
 	for i := 0; i < runtime.NumCPU(); i++ {
 		l, err := lc.ListenPacket(context.Background(), "udp", address)
 		if err != nil {

wpex.go

@@ -4,12 +4,13 @@ import (
 	"encoding/base64"
 	"flag"
 	"fmt"
-	"github.com/weiiwang01/wpex/internal/relay"
-	"golang.org/x/time/rate"
-	"log"
 	"log/slog"
 	"os"
 	"strings"
+
+	"github.com/weiiwang01/wpex/internal/relay"
+	"golang.org/x/time/rate"
+	"gopkg.in/ini.v1"
 )
 
 var version string
@@ -31,8 +32,9 @@ func main() {
 	debug := flag.Bool("debug", false, "enable debug messages")
 	broadcastRate := flag.Uint("broadcast-rate", 0, "broadcast rate limit in packet per second")
 	versionFlag := flag.Bool("version", false, "show version number and quit")
-	var allows pubKeys
-	flag.Var(&allows, "allow", "allow a wireguard public key. --allow can be used multiple times for allowing multiple public keys")
+	var allowsFromFlags pubKeys
+	flag.Var(&allowsFromFlags, "allow", "allow a wireguard public key. --allow can be used multiple times for allowing multiple public keys")
+	allowFromConfPath := flag.String("allow-from-conf", "", "allow public keys from a wireguard configuration file.")
 	flag.Parse()
 	if *versionFlag {
 		fmt.Println("wpex", version)
@@ -45,21 +47,63 @@ func main() {
 	}
 	slog.SetDefault(logger)
 	address := fmt.Sprintf("%s:%d", *bind, *port)
-	var allowKeys [][]byte
-	for _, allow := range allows {
-		k, err := base64.StdEncoding.DecodeString(allow)
-		if err != nil || len(k) != 32 {
-			log.Fatal(fmt.Sprintf("invalid wireguard public key: '%s'", allow))
-		}
-		logger.Debug("allow wireguard public key", "key", allow)
-		allowKeys = append(allowKeys, k)
-	}
+
 	limit := rate.Limit(*broadcastRate)
 	if *broadcastRate == 0 {
 		slog.Debug("broadcast rate limit is set to +Inf")
 		limit = rate.Inf
 	} else {
 		slog.Debug(fmt.Sprintf("broadcast rate limit is set to %d", *broadcastRate))
 	}
-	relay.Start(address, allowKeys, rate.NewLimiter(limit, int((*broadcastRate)*5)))
+
+	getAllAllowKeys := func() ([][]byte, error) {
+		var allows pubKeys
+
+		// add keys from flags
+		for _, allowFromFlag := range allowsFromFlags {
+			allows = append(allows, allowFromFlag)
+		}
+
+		// add keys from conf
+		if *allowFromConfPath != "" {
+			logger.Debug("allow wireguard public keys from configuration file", "file", *allowFromConfPath)
+			cfg, err := ini.LoadSources(
+				ini.LoadOptions{
+					AllowNonUniqueSections: true,
+					InsensitiveKeys:        true,
+				},
+				*allowFromConfPath,
+			)
+			if err != nil {
+				return nil, err
+			}
+			for _, section := range cfg.Sections() {
+				if strings.EqualFold(section.Name(), "Peer") && section.HasKey("PublicKey") {
+					allowFromConf := section.Key("PublicKey").String()
+					allows = append(allows, allowFromConf)
+				}
+			}
+		}
+
+		var invalidKeysFound bool
+		var allowKeys [][]byte
+		for _, allow := range allows {
+			k, err := base64.StdEncoding.DecodeString(allow)
+			if err != nil || len(k) != 32 {
+				slog.Error(fmt.Sprintf("invalid wireguard public key: '%s'", allow))
+				invalidKeysFound = true
+			} else {
+				logger.Debug("allow wireguard public key", "key", allow)
+				allowKeys = append(allowKeys, k)
+			}
+		}
+
+		if invalidKeysFound {
+			return nil, fmt.Errorf("invalid wireguard public keys found")
+		}
+
+		return allowKeys, nil
+	}
+
+	relay.Start(address, getAllAllowKeys, rate.NewLimiter(limit, int((*broadcastRate)*5)))
 }