From 909852b36fe6fdd0b8c09128403bb1615b608bb6 Mon Sep 17 00:00:00 2001
From: NHAS <jordanatararimu@gmail.com>
Date: Mon, 10 Jun 2024 19:01:02 +1200
Subject: [PATCH] Add reset, add mising boolean

---
 internal/router/session_manager.go             | 18 ++++++++++++++++++
 .../webserver/resources/static/js/challenge.js | 14 +++++++++++---
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/internal/router/session_manager.go b/internal/router/session_manager.go
index 591d8b2..d09e983 100644
--- a/internal/router/session_manager.go
+++ b/internal/router/session_manager.go
@@ -33,6 +33,8 @@ func (ws *wsConnWrapper) Close() error {
 		return nil
 	}
 
+	ws.isClosed = true
+
 	close(ws.wait)
 	return ws.Conn.Close()
 }
@@ -74,6 +76,9 @@ func (c *Challenger) Challenge(address string) error {
 	c.RLock()
 	defer c.RUnlock()
 
+	log.Println("challenging")
+	defer log.Println("end challenging")
+
 	conn, ok := c.connections[address]
 	if !ok {
 		return fmt.Errorf("no connection found for device: %s", address)
@@ -125,6 +130,18 @@ func (c *Challenger) Challenge(address string) error {
 	return nil
 }
 
+func (c *Challenger) Reset(address string) {
+	c.RLock()
+	defer c.RUnlock()
+
+	conn, ok := c.connections[address]
+	if !ok {
+		return
+	}
+
+	conn.WriteJSON("reset")
+}
+
 func (c *Challenger) WS(w http.ResponseWriter, r *http.Request) {
 	remoteAddress := utils.GetIPFromRequest(r)
 	user, err := users.GetUserFromAddress(remoteAddress)
@@ -166,6 +183,7 @@ func (c *Challenger) WS(w http.ResponseWriter, r *http.Request) {
 
 	err = c.Challenge(remoteAddress.String())
 	if err != nil {
+		c.Reset(remoteAddress.String())
 		log.Printf("client did not complete inital ws challenge: %s", err)
 		return
 	}
diff --git a/internal/webserver/resources/static/js/challenge.js b/internal/webserver/resources/static/js/challenge.js
index 8a77863..64bdce1 100644
--- a/internal/webserver/resources/static/js/challenge.js
+++ b/internal/webserver/resources/static/js/challenge.js
@@ -23,9 +23,17 @@ function connect() {
 
     ws.onmessage = function (e) {
         console.log('Message:', e.data);
-        ws.send(
-            JSON.stringify({challenge: challenge
-        }));
+        switch(e.data) {
+            case "challenge":
+                ws.send(
+                    JSON.stringify({challenge: challenge
+                }));
+            return
+            case "reset":
+                window.location.href = '/'
+            return
+        }
+   
     };
 
     ws.onclose = function (e) {