From 0cbf5e431c0af74b3de9dbe65e2ebc46050620fe Mon Sep 17 00:00:00 2001
From: mkwcat <cat@mkw.re>
Date: Wed, 31 Jan 2024 22:02:54 -0500
Subject: [PATCH] QR2: Clear message ack waker before adding to sleeper

---
 qr2/main.go    | 9 ++++++++-
 qr2/message.go | 4 +++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/qr2/main.go b/qr2/main.go
index f0bfb59..2383fcf 100644
--- a/qr2/main.go
+++ b/qr2/main.go
@@ -120,6 +120,9 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 	case ClientMessageAckRequest:
 		logging.Info(moduleName, "Command:", aurora.Yellow("CLIENT_MESSAGE_ACK"))
 
+		mutex.Lock()
+		defer mutex.Unlock()
+
 		// In case ClientExploitReply is lost, this can be checked as well
 		// This would be sent either after the payload is downloaded, or the client is already patched
 		session.ExploitReceived = true
@@ -135,8 +138,9 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 		conn.WriteTo(createResponseHeader(KeepAliveRequest, 0), addr)
 
 		mutex.Lock()
+		defer mutex.Unlock()
+
 		session.LastKeepAlive = time.Now().Unix()
-		mutex.Unlock()
 		return
 
 	case AvailableRequest:
@@ -151,6 +155,9 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 	case ClientExploitReply:
 		logging.Info(moduleName, "Command:", aurora.Yellow("CLIENT_EXPLOIT_ACK"))
 
+		mutex.Lock()
+		defer mutex.Unlock()
+
 		session.ExploitReceived = true
 		if login := session.Login; login != nil {
 			login.NeedsExploit = false
diff --git a/qr2/message.go b/qr2/message.go
index 348946f..360c838 100644
--- a/qr2/message.go
+++ b/qr2/message.go
@@ -241,11 +241,13 @@ func SendClientMessage(senderIP string, destSearchID uint64, message []byte) {
 	}
 
 	s := sleep.Sleeper{}
+
+	receiver.MessageAckWaker.Clear()
 	s.AddWaker(receiver.MessageAckWaker)
+
 	timeWaker := sleep.Waker{}
 	s.AddWaker(&timeWaker)
 
-	receiver.MessageAckWaker.Clear()
 	timeOutCount := 0
 	for {
 		time.AfterFunc(1*time.Second, func() {