From 00ffe320024d3e740cf507929f8eaffca8c16ca0 Mon Sep 17 00:00:00 2001
From: MikeIsAStar <99037623+MikeIsAStar@users.noreply.github.com>
Date: Sat, 9 Dec 2023 17:55:00 -0500
Subject: [PATCH] QR2: Prevent invalid packets from stopping the server

---
 qr2/main.go | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/qr2/main.go b/qr2/main.go
index 21d8614..31867cb 100644
--- a/qr2/main.go
+++ b/qr2/main.go
@@ -54,16 +54,19 @@ func StartServer() {
 }
 
 func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
-	if buffer[0] == AvailableRequest {
-		logging.Notice("QR2", "Command:", aurora.Yellow("AVAILABLE"))
-		conn.WriteTo(createResponseHeader(AvailableRequest, 0), addr)
-		return
-	}
-
+	packetType := buffer[0]
 	sessionId := binary.BigEndian.Uint32(buffer[1:5])
+	session, ok := sessions[sessionId]
 	moduleName := "QR2:" + strconv.FormatInt(int64(sessionId), 10)
 
-	switch buffer[0] {
+	if packetType != HeartbeatRequest && packetType != AvailableRequest {
+		if !ok {
+			logging.Error(moduleName, "Invalid session")
+			return
+		}
+	}
+
+	switch packetType {
 	case QueryRequest:
 		logging.Notice(moduleName, "Command:", aurora.Yellow("QUERY"))
 		break
@@ -72,9 +75,9 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 		logging.Notice(moduleName, "Command:", aurora.Yellow("CHALLENGE"))
 
 		mutex.Lock()
-		if sessions[sessionId].Challenge != "" {
+		if session.Challenge != "" {
 			// TODO: Verify the challenge
-			sessions[sessionId].Authenticated = true
+			session.Authenticated = true
 			mutex.Unlock()
 
 			conn.WriteTo(createResponseHeader(ClientRegisteredReply, sessionId), addr)
@@ -110,13 +113,14 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 
 	case KeepAliveRequest:
 		logging.Notice(moduleName, "Command:", aurora.Yellow("KEEPALIVE"))
-		sessionId := binary.BigEndian.Uint32(buffer[1:5])
 		mutex.Lock()
-		sessions[sessionId].LastKeepAlive = time.Now().Unix()
+		session.LastKeepAlive = time.Now().Unix()
 		mutex.Unlock()
 		return
 
 	case AvailableRequest:
+		logging.Notice("QR2", "Command:", aurora.Yellow("AVAILABLE"))
+		conn.WriteTo(createResponseHeader(AvailableRequest, 0), addr)
 		return
 
 	case ClientRegisteredReply: