From 08b7446dd90d2760ebe7e908511eb4c2680ebcac Mon Sep 17 00:00:00 2001
From: NHAS <jordanatararimu@gmail.com>
Date: Wed, 8 May 2024 11:37:52 +1200
Subject: [PATCH] Fix issue where just after startup wag would be unable to
 deauthenticate all user devices

---
 internal/router/bpf.go          |  1 +
 internal/router/statemachine.go |  6 ++++--
 internal/router/wireguard.go    | 12 ++++++++++++
 internal/users/user.go          |  1 -
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/internal/router/bpf.go b/internal/router/bpf.go
index e749aa25..e6053428 100644
--- a/internal/router/bpf.go
+++ b/internal/router/bpf.go
@@ -163,6 +163,7 @@ func setupXDP(users []data.UserModel, knownDevices []data.Device) error {
 	}
 
 	for _, device := range knownDevices {
+
 		err := xdpAddDevice(device.Username, device.Address)
 		if err != nil {
 			return errors.New("xdp setup add device to user: " + err.Error())
diff --git a/internal/router/statemachine.go b/internal/router/statemachine.go
index 01957010..6983d230 100644
--- a/internal/router/statemachine.go
+++ b/internal/router/statemachine.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/NHAS/wag/internal/acls"
 	"github.com/NHAS/wag/internal/data"
+	"github.com/NHAS/wag/internal/webserver/authenticators/types"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
 
@@ -136,7 +137,7 @@ func userChanges(key string, current data.UserModel, previous data.UserModel, et
 		}
 	case data.MODIFIED:
 
-		if current.Locked != previous.Locked {
+		if current.Locked != previous.Locked || current.Locked {
 
 			lock := uint32(1)
 			if !current.Locked {
@@ -150,7 +151,8 @@ func userChanges(key string, current data.UserModel, previous data.UserModel, et
 			}
 		}
 
-		if current.Mfa != previous.Mfa || current.MfaType != previous.MfaType {
+		if current.Mfa != previous.Mfa || current.MfaType != previous.MfaType ||
+			!current.Enforcing || types.MFA(current.MfaType) == types.Unset {
 			err := DeauthenticateAllDevices(current.Username)
 			if err != nil {
 				log.Printf("cannot deauthenticate user %s: %s", current.Username, err)
diff --git a/internal/router/wireguard.go b/internal/router/wireguard.go
index cef6e63a..4e58da70 100644
--- a/internal/router/wireguard.go
+++ b/internal/router/wireguard.go
@@ -47,6 +47,8 @@ func (msg *IfAddrmsg) Serialize() []byte {
 }
 
 func setupWireguard(devices []data.Device) error {
+	lock.Lock()
+	defer lock.Unlock()
 
 	var c wgtypes.Config
 
@@ -102,6 +104,16 @@ func setupWireguard(devices []data.Device) error {
 			d := time.Duration(config.Values.Wireguard.ServerPersistentKeepAlive) * time.Second
 			pc.PersistentKeepaliveInterval = &d
 		}
+
+		addressesMap, ok := usersToAddresses[device.Username]
+		if !ok {
+			addressesMap = make(map[string]string)
+		}
+
+		addressesMap[device.Address] = pk.String()
+		usersToAddresses[device.Username] = addressesMap
+		addressesToUsers[device.Address] = device.Username
+
 		c.Peers = append(c.Peers, pc)
 	}
 
diff --git a/internal/users/user.go b/internal/users/user.go
index 8de44e5b..a5e1f5a4 100644
--- a/internal/users/user.go
+++ b/internal/users/user.go
@@ -22,7 +22,6 @@ func (u *user) ResetDeviceAuthAttempts(address string) error {
 
 func (u *user) ResetMfa() error {
 
-	// the MFA column is marked as "unique" so just set it as the username as that is also unique
 	err := data.SetUserMfa(u.Username, u.Username, string(types.Unset))
 	if err != nil {
 		return err