From b9c063c41e291e2b1ecf9b37f5352f37006e0051 Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Tue, 5 Mar 2024 14:47:02 -0500
Subject: [PATCH] [O] Reject unauthenticated aimedb requests

---
 .../icu/samnyan/aqua/net/db/AquaNetUser.kt     | 16 ++++++++++++----
 .../icu/samnyan/aqua/sega/aimedb/AimeDB.kt     | 18 +++++++++++-------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/main/java/icu/samnyan/aqua/net/db/AquaNetUser.kt b/src/main/java/icu/samnyan/aqua/net/db/AquaNetUser.kt
index 525d6f09..bb4fea77 100644
--- a/src/main/java/icu/samnyan/aqua/net/db/AquaNetUser.kt
+++ b/src/main/java/icu/samnyan/aqua/net/db/AquaNetUser.kt
@@ -1,10 +1,9 @@
 package icu.samnyan.aqua.net.db
 
 import com.fasterxml.jackson.annotation.JsonIgnore
-import ext.Str
-import ext.async
-import ext.isValidEmail
-import ext.minus
+import ext.*
+import icu.samnyan.aqua.sega.allnet.AllNetProps
+import icu.samnyan.aqua.sega.allnet.KeyChipRepo
 import icu.samnyan.aqua.sega.allnet.KeychipSession
 import icu.samnyan.aqua.sega.general.dao.CardRepository
 import icu.samnyan.aqua.sega.general.model.Card
@@ -101,6 +100,8 @@ class AquaUserServices(
     val userRepo: AquaNetUserRepo,
     val cardRepo: CardRepository,
     val hasher: PasswordEncoder,
+    val keyChipRepo: KeyChipRepo,
+    val allNetProps: AllNetProps
 ) {
     companion object {
         val SETTING_FIELDS = AquaUserServices::class.functions
@@ -121,6 +122,13 @@ class AquaUserServices(
             ?.let { callback(it) } ?: (404 - "Card not found")
         else byName(username) { callback(it.ghostCard) }
 
+    fun validKeychip(keychipId: Str): Bool {
+        if (!allNetProps.checkKeychip) return true
+        if (keychipId.isBlank()) return false
+        if (userRepo.findByKeychip(keychipId) != null || keyChipRepo.existsByKeychipId(keychipId)) return true
+        return false
+    }
+
     fun checkUsername(username: Str) = username.apply {
         // Check if username is valid
         if (length < 2) 400 - "Username must be at least 2 letters"
diff --git a/src/main/java/icu/samnyan/aqua/sega/aimedb/AimeDB.kt b/src/main/java/icu/samnyan/aqua/sega/aimedb/AimeDB.kt
index eadae568..bffe425c 100644
--- a/src/main/java/icu/samnyan/aqua/sega/aimedb/AimeDB.kt
+++ b/src/main/java/icu/samnyan/aqua/sega/aimedb/AimeDB.kt
@@ -1,6 +1,7 @@
 package icu.samnyan.aqua.sega.aimedb
 
 import ext.toHex
+import icu.samnyan.aqua.net.db.AquaUserServices
 import icu.samnyan.aqua.sega.general.model.Card
 import icu.samnyan.aqua.sega.general.service.CardService
 import io.netty.buffer.ByteBuf
@@ -22,7 +23,8 @@ import kotlin.jvm.optionals.getOrNull
 @Component
 @ChannelHandler.Sharable
 class AimeDB(
-    val cardService: CardService
+    val cardService: CardService,
+    val us: AquaUserServices,
 ): ChannelInboundHandlerAdapter() {
     val logger: Logger = LoggerFactory.getLogger(AimeDB::class.java)
 
@@ -53,20 +55,22 @@ class AimeDB(
      * Handle the incoming request
      */
     override fun channelRead(ctx: ChannelHandlerContext, msg: Any) {
-        if (msg is Map<*, *>) {
+        if (msg !is Map<*, *>) return
+        try {
             val type = msg["type"] as Int
             val data = msg["data"] as ByteBuf
             val base = getBaseInfo(data)
-            val handler = handlers[type] ?: let {
-                logger.error("AimeDB: Unknown request type 0x${type.toString(16)}")
-                ctx.flush()
-                return
-            }
+            val handler = handlers[type] ?: return logger.error("AimeDB: Unknown request type 0x${type.toString(16)}")
 
             logger.info("AimeDB /${handler.name} : (game ${base.gameId}, keychip ${base.keychipId})")
 
+            // Check keychip
+            if (!us.validKeychip(base.keychipId)) return logger.warn("> Rejected: Keychip not found")
+
             handler.fn(data)?.let { ctx.write(it) }
+        } finally {
             ctx.flush()
+            ctx.close()
         }
     }