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() {