QR2: Send QR2 exploit immediately for Mario Kart DS

common/game_list.go

@@ -102,3 +102,8 @@ func GetExpectedUnitCode(gameName string) byte {
 
 	return 0
 }
+
+func DoesGameNeedExploit(gameName string) bool {
+	// Exploit is only implemented for Mario Kart Wii and Mario Kart DS currently
+	return gameName == "mariokartwii" || gameName == "mariokartds"
+}

gpcm/login.go

@@ -183,8 +183,17 @@ func (g *GameSpySession) login(command common.GameSpyCommand) {
 
 	if g.UnitCode == UnitCodeWii {
 		if isLocalhost && !payloadVerExists && !signatureExists {
-			// Players using the DNS exploit, need patching using a QR2 exploit
-			// TODO: Check that the game is compatible with the DNS
+			// Players using the DNS, need patching using a QR2 exploit
+			if !common.DoesGameNeedExploit(g.GameName) {
+				logging.Error(g.ModuleName, "Using DNS for incompatible game:", aurora.Cyan(g.GameName))
+				g.replyError(GPError{
+					ErrorCode:   ErrLogin.ErrorCode,
+					ErrorString: "The client is not patched to use WiiLink WFC.",
+					Fatal:       true,
+				})
+				return
+			}
+
 			g.NeedsExploit = true
 		} else {
 			deviceId = g.verifyExLoginInfo(command, authToken)
@@ -192,6 +201,12 @@ func (g *GameSpySession) login(command common.GameSpyCommand) {
 				return
 			}
 		}
+	} else if g.UnitCode == UnitCodeDS {
+		g.NeedsExploit = common.DoesGameNeedExploit(g.GameName)
+	} else {
+		logging.Error(g.ModuleName, "Invalid unit code specified:", aurora.Cyan(unitcd))
+		g.replyError(ErrLogin)
+		return
 	}
 
 	response := generateResponse(g.Challenge, challenge, authToken, command.OtherValues["challenge"])

nas/auth.go

@@ -113,6 +113,10 @@ func handleAuthRequest(moduleName string, w http.ResponseWriter, r *http.Request
 			reply = svcloc(fields)
 			break
 
+		case "SVCLOC":
+			reply = svcloc(fields)
+			break
+
 		default:
 			logging.Error(moduleName, "Unknown action:", aurora.Cyan(action))
 			reply = map[string]string{

qr2/heartbeat.go

@@ -96,27 +96,33 @@ func heartbeat(moduleName string, conn net.PacketConn, addr net.Addr, buffer []b
 		return
 	}
 
-	if payload["gamename"] == "mariokartwii" && len(unknowns) > 0 {
+	if len(unknowns) > 0 {
 		// Try to login using the first unknown as a profile ID
 		// This makes it possible to execute the exploit on the client sooner
 
 		mutex.Lock()
-		session, sessionExists := sessions[lookupAddr]
+		sessionPtr, sessionExists := sessions[lookupAddr]
 		if !sessionExists {
 			logging.Error(moduleName, "Session not found")
-		} else if session.Login == nil {
+		} else if sessionPtr.Login == nil {
 			profileId := unknowns[0]
 			logging.Info(moduleName, "Attempting to use unknown as profile ID", aurora.Cyan(profileId))
-			session.setProfileID(moduleName, profileId)
+			sessionPtr.setProfileID(moduleName, profileId)
 		}
+		session = *sessionPtr
 		mutex.Unlock()
 	}
 
 	if !session.Authenticated || noIP {
 		sendChallenge(conn, addr, session, lookupAddr)
-	} else if !session.ExploitReceived && session.Login != nil && session.Login.NeedsExploit && statechanged == "1" {
-		logging.Notice(moduleName, "Sending SBCM exploit to DNS patcher client")
-		sendClientExploit(moduleName, session)
+	}
+
+	if !session.ExploitReceived && session.Login != nil && session.Login.NeedsExploit {
+		// The version of DWC in Mario Kart DS doesn't check matching status
+		if (!noIP && statechanged == "1") || payload["gamename"] == "mariokartds" {
+			logging.Notice(moduleName, "Sending SBCM exploit to DNS patcher client")
+			sendClientExploit(moduleName, session)
+		}
 	}
 
 	mutex.Lock()

qr2/main.go

@@ -123,6 +123,10 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 		// 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
+		if login := session.Login; login != nil {
+			login.NeedsExploit = false
+		}
+
 		session.MessageAckWaker.Assert()
 		return