From d7e727df0b4a5cd144cb36f6172ec91bb1406fa0 Mon Sep 17 00:00:00 2001
From: NHAS <jordanatararimu@gmail.com>
Date: Mon, 10 Jun 2024 18:33:51 +1200
Subject: [PATCH] Allow for 3 seconds for grace before challenge is fatal

---
 internal/router/session_manager.go |  6 +++++-
 internal/router/statemachine.go    | 14 ++++++++++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/internal/router/session_manager.go b/internal/router/session_manager.go
index b1432ca..960b391 100644
--- a/internal/router/session_manager.go
+++ b/internal/router/session_manager.go
@@ -81,22 +81,26 @@ func (c *Challenger) Challenge(address string) error {
 
 	err := conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
 	if err != nil {
+		conn.Close()
 		return err
 	}
 
 	err = conn.WriteJSON("challenge")
 	if err != nil {
+		conn.Close()
 		return err
 	}
 
 	err = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
 	if err != nil {
+		conn.Close()
 		return err
 	}
 
 	msg := struct{ Challenge string }{}
 	err = conn.ReadJSON(&msg)
 	if err != nil {
+		conn.Close()
 		return err
 	}
 
@@ -153,7 +157,7 @@ func (c *Challenger) WS(w http.ResponseWriter, r *http.Request) {
 
 	err = c.Challenge(remoteAddress.String())
 	if err != nil {
-		log.Printf("client did not complete ws challenge: %s", err)
+		log.Printf("client did not complete inital ws challenge: %s", err)
 		return
 	}
 
diff --git a/internal/router/statemachine.go b/internal/router/statemachine.go
index 62d7f7c..1af03b6 100644
--- a/internal/router/statemachine.go
+++ b/internal/router/statemachine.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log"
 	"strings"
+	"time"
 
 	"github.com/NHAS/wag/internal/acls"
 	"github.com/NHAS/wag/internal/data"
@@ -103,8 +104,16 @@ func deviceChanges(_ string, current, previous data.Device, et data.EventType) e
 
 			log.Printf("challenging %s:%s device, as endpoint changed: %s -> %s", current.Username, current.Address, current.Endpoint.String(), previous.Endpoint.String())
 			// Will take at most 4 seconds
-			err := Verifier.Challenge(current.Address)
-			if err != nil {
+
+			attempts := 0
+			for ; attempts < 3; attempts++ {
+				err = Verifier.Challenge(current.Address)
+				if err != nil {
+					time.Sleep(1 * time.Second)
+				}
+			}
+
+			if attempts >= 3 {
 				log.Printf("%s:%s failed to pass websockets challenge: %s", current.Username, current.Address, err)
 				err := Deauthenticate(current.Address)
 				if err != nil {
@@ -112,6 +121,7 @@ func deviceChanges(_ string, current, previous data.Device, et data.EventType) e
 				}
 			} else {
 				log.Printf("%s:%s device succeeded challenge", current.Username, current.Address)
+
 			}
 		}