feat: support multiple keychips per user (#215)

Co-authored-by: hykilpikonna <22280294+hykilpikonna@users.noreply.github.com>
Co-authored-by: asterisk727 <59166650+asterisk727@users.noreply.github.com>
Co-authored-by: Raymond <101374892+raymonable@users.noreply.github.com>

src/main/java/icu/samnyan/aqua/net/Bot.kt

@@ -116,7 +116,7 @@ class BotController(
             val gamesDict = listOfNotNull(chu3, mai2).map {
                 // Find the keychip owner
                 val keychip = it.lastClientId
-                val keychipOwner = keychip?.let { us.userRepo.findByKeychip(it) }
+                val keychipOwner = keychip?.let { k -> us.userKeychipRepo.findByKeychipId(k)?.user }
 
                 mapOf(
                     "userData" to it,

src/main/java/icu/samnyan/aqua/net/UserRegistrar.kt

@@ -2,12 +2,16 @@ package icu.samnyan.aqua.net
 
 import ext.*
 import icu.samnyan.aqua.net.components.*
+import icu.samnyan.aqua.net.db.AquaNetUser
 import icu.samnyan.aqua.net.db.AquaNetUserRepo
 import icu.samnyan.aqua.net.db.AquaUserServices
 import icu.samnyan.aqua.net.db.EmailConfirmationRepo
 import icu.samnyan.aqua.net.db.ResetPasswordRepo
+import icu.samnyan.aqua.net.utils.AquaNetProps
 import icu.samnyan.aqua.net.utils.PathProps
 import icu.samnyan.aqua.net.utils.SUCCESS
+import icu.samnyan.aqua.sega.allnet.UserKeychip
+import icu.samnyan.aqua.sega.allnet.UserKeychipRepo
 import icu.samnyan.aqua.sega.general.dao.CardRepository
 import icu.samnyan.aqua.sega.general.model.CardStatus
 import jakarta.servlet.http.HttpServletRequest
@@ -33,6 +37,8 @@ class UserRegistrar(
     val cardRepo: CardRepository,
     val validator: AquaUserServices,
     val emailProps: EmailProperties,
+    val userKeychipRepo: UserKeychipRepo,
+    val aquaNetProps: AquaNetProps,
     final val paths: PathProps
 ) {
     val portraitPath = paths.aquaNetPortrait.path()
@@ -235,20 +241,57 @@ class UserRegistrar(
         SUCCESS
     }
 
+    val keychipPattern = Regex("^A\\d{14}$")
+    val dashedKeychipPattern = Regex("^A\\d{3}-\\d{11}$")
     val keychipRange = 1e9.toULong()..1e10.toULong() - 1UL
 
+
+    private fun ensureCanModifyKeychips(u: AquaNetUser) {
+        if (!u.canModifyKeychips) 403 - "You don't have permission to modify keychips"
+    }
+
+    private fun validateCustomKeychip(keychipId: Str): Str {
+        val raw = keychipId.trim().uppercase()
+        val normalized = raw.replace("-", "")
+
+        val validRawFormat = raw == normalized || dashedKeychipPattern.matches(raw)
+        if (!validRawFormat || !keychipPattern.matches(normalized))
+            400 - "Invalid keychip format. Expected A followed by 14 digits (with optional dash)"
+
+        return normalized
+    }
+
     @API("/keychip")
-    @Doc("Get a Keychip ID so that the user can connect to the server.", "Success message")
-    suspend fun setupConnection(@RP token: Str) = jwt.auth(token) { u ->
-        u.keychip?.let { return mapOf("keychip" to it) }
-        log.info("Net: /user/keychip setup: ${u.auId} for ${u.username}")
+    @Doc("List all keychip IDs associated with the current user's account.", "List of keychip IDs")
+    suspend fun listKeychips(@RP token: Str) = jwt.auth(token) { u ->
+        val keychips = async { userKeychipRepo.findAllByUserAuId(u.auId) }
+        mapOf("keychips" to keychips.map { it.keychipId })
+    }
 
-        // Generate a keychip id with 10 digits (e.g. A1234567890)
-        var new = "A" + keychipRange.random()
-        while (async { userRepo.findByKeychip(new) != null }) new = "A" + keychipRange.random()
-        async { userRepo.save(u.apply { keychip = new }) }
+    @API("/keychip/add")
+    @Doc("Add a custom keychip ID for the user.", "The newly added keychip ID")
+    suspend fun addKeychip(@RP token: Str, @RP keychipId: Str) = jwt.auth(token) { u ->
+        ensureCanModifyKeychips(u)
+        log.info("Net: /user/keychip/add: ${u.auId} for ${u.username}")
+        val validated = validateCustomKeychip(keychipId)
 
-        mapOf("keychip" to new)
+        if (async { userKeychipRepo.existsByKeychipId(validated) })
+            400 - "Keychip already exists"
+
+        if (userKeychipRepo.findAllByUserAuId(u.auId).size >= aquaNetProps.keychipLimit)
+            400 - "Exceeds maximum keychip count"
+
+        async { userKeychipRepo.save(UserKeychip(user = u, keychipId = validated)) }
+        mapOf("keychipId" to validated)
+    }
+
+    @API("/keychip/delete")
+    @Doc("Remove a keychip ID from the user's account.", "Success message")
+    suspend fun deleteKeychip(@RP token: Str, @RP keychipId: Str) = jwt.auth(token) { u ->
+        ensureCanModifyKeychips(u)
+        val deleted = async { userKeychipRepo.deleteByKeychipIdAndUserAuId(keychipId, u.auId) }
+        if (deleted == 0L) 404 - "Keychip not found"
+        SUCCESS
     }
 
     @API("/upload-pfp", consumes = ["multipart/form-data"])

src/main/java/icu/samnyan/aqua/net/db/AquaNetUser.kt

@@ -6,8 +6,9 @@ import icu.samnyan.aqua.net.UserRegistrar.Companion.cardExtIdEnd
 import icu.samnyan.aqua.net.UserRegistrar.Companion.cardExtIdStart
 import icu.samnyan.aqua.net.components.JWT
 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.allnet.UserKeychip
+import icu.samnyan.aqua.sega.allnet.UserKeychipRepo
 import icu.samnyan.aqua.sega.general.GameMusicPopularity
 import icu.samnyan.aqua.sega.general.dao.CardRepository
 import icu.samnyan.aqua.sega.general.model.Card
@@ -74,16 +75,19 @@ class AquaNetUser(
     @OneToMany(mappedBy = "aquaUser", cascade = [CascadeType.ALL])
     var cards: MutableList<Card> = mutableListOf(),
 
-    // Each user can have one keychip (if the user owns a cabinet)
+    // Each user can have multiple keychips (if the user owns cabinets)
     @JsonIgnore
-    @Column(nullable = true, length = 32, unique = true)
-    var keychip: Str? = null,
+    @OneToMany(mappedBy = "user", cascade = [CascadeType.ALL])
+    var keychips: MutableList<UserKeychip> = mutableListOf(),
 
     // Each user's keychip can have multiple sessions
     @JsonIgnore
     @OneToMany(mappedBy = "user", cascade = [CascadeType.ALL])
     var keychipSessions: MutableList<KeychipSession> = mutableListOf(),
 
+    @Column(nullable = false)
+    var canModifyKeychips: Boolean = false,
+
     @OneToOne(cascade = [CascadeType.ALL])
     @JoinColumn(name = "gameOptions", unique = true, nullable = true)
     var gameOptions: AquaGameOptions? = null,
@@ -105,7 +109,6 @@ interface AquaNetUserRepo : JpaRepository<AquaNetUser, Long> {
     fun findByAuId(auId: Long): AquaNetUser?
     fun findByEmailIgnoreCase(email: String): AquaNetUser?
     fun findByUsernameIgnoreCase(username: String): AquaNetUser?
-    fun findByKeychip(keychip: String): AquaNetUser?
     fun findByGhostCardExtId(extId: Long): AquaNetUser?
 }
 
@@ -124,7 +127,7 @@ class AquaUserServices(
     val userRepo: AquaNetUserRepo,
     val cardRepo: CardRepository,
     val hasher: PasswordEncoder,
-    val keyChipRepo: KeyChipRepo,
+    val userKeychipRepo: UserKeychipRepo,
     val allNetProps: AllNetProps,
     val jwt: JWT,
     val em: EntityManager,
@@ -142,9 +145,18 @@ class AquaUserServices(
             }
     }
 
+    val keychipRange = 1e9.toULong()..1e10.toULong() - 1UL
+    private fun generateKeychipId(): String {
+        // 1337 is retained for backwards compatibility, it's no longer required for cabinet keychips
+        var keychip = "A" + keychipRange.random() + "1337"
+        while ( userKeychipRepo.existsByKeychipId(keychip) )
+            keychip = "A" + keychipRange.random() + "1337"
+        return keychip
+    }
+
     fun create(username: Str, email: Str, password: Str, country: Str, emailConfirmed: Boolean = false): AquaNetUser {
         // Create user
-        val u = AquaNetUser(
+        val user = AquaNetUser(
             username = checkUsername(username),
             email = validateEmail(email),
             pwHash = checkPwHash(password),
@@ -158,16 +170,20 @@ class AquaUserServices(
             luid = extId.toString()
             registerTime = LocalDateTime.now()
             accessTime = registerTime
-            aquaUser = u
+            aquaUser = user
             isGhost = true
         }
-        u.ghostCard = card
+        user.ghostCard = card
+
+        // Create an automatic keychip
+        val keychip = UserKeychip(0, user, generateKeychipId())
 
         // Save the user
-        userRepo.save(u)
+        userRepo.save(user)
         cardRepo.save(card)
+        userKeychipRepo.save(keychip)
 
-        return u
+        return user
     }
 
     fun update(user: AquaNetUser, key: Str, value: Str) {
@@ -192,7 +208,7 @@ class AquaUserServices(
     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
+        if (userKeychipRepo.existsByKeychipId(keychipId)) return true
         return false
     }
 

src/main/java/icu/samnyan/aqua/net/games/GameApiController.kt

@@ -4,6 +4,7 @@ import ext.*
 import icu.samnyan.aqua.net.BotProps
 import icu.samnyan.aqua.net.db.AquaUserServices
 import icu.samnyan.aqua.net.utils.SUCCESS
+import icu.samnyan.aqua.sega.allnet.UserKeychipRepo
 import icu.samnyan.aqua.sega.general.model.Card
 import icu.samnyan.aqua.sega.general.service.CardService
 import jakarta.annotation.PostConstruct
@@ -32,6 +33,7 @@ abstract class GameApiController<T : IUserData>(val name: String, userDataClass:
     open val gettableFields: Set<String> = setOf()
 
     @Autowired lateinit var cardService: CardService
+    @Autowired lateinit var userKeychipRepo: UserKeychipRepo
 
     @API("trend")
     abstract suspend fun trend(@RP username: String): List<TrendOut>
@@ -205,7 +207,7 @@ abstract class GameApiController<T : IUserData>(val name: String, userDataClass:
             lastVersion = user.lastRomVersion,
             ratingComposition = ratingComp,
             recent = plays.sortedBy { it.userPlayDate.toString() }.takeLast(100).reversed(),
-            lastPlayedHost = user.lastClientId?.let { us.userRepo.findByKeychip(it)?.username },
+            lastPlayedHost = user.lastClientId?.let { userKeychipRepo.findByKeychipId(it)?.user?.username },
             rival = rival,
             favorites = favorites
         )

src/main/java/icu/samnyan/aqua/net/utils/AquaNetProps.kt

@@ -10,6 +10,7 @@ import kotlin.io.path.createDirectories
 @ConfigurationProperties(prefix = "aqua-net")
 class AquaNetProps {
     var linkCardLimit: Int = 10
+    var keychipLimit: Int = 20
     var importBackupPath = "data/import-backups"
 
     @PostConstruct

src/main/java/icu/samnyan/aqua/sega/allnet/AllNet.kt

@@ -1,7 +1,6 @@
 package icu.samnyan.aqua.sega.allnet
 
 import ext.*
-import icu.samnyan.aqua.net.db.AquaNetUserRepo
 import icu.samnyan.aqua.sega.allnet.AllNetBillingDecoder.decodeAllNet
 import jakarta.servlet.http.HttpServletRequest
 import jakarta.servlet.http.HttpServletResponse
@@ -59,9 +58,8 @@ class AllNetProps {
 @Suppress("HttpUrlsUsage")
 @RestController
 class AllNet(
-    val userRepo: AquaNetUserRepo,
+    val userKeychipRepo: UserKeychipRepo,
     val keychipSessionService: KeychipSessionService,
-    val keychipRepo: KeyChipRepo,
     val props: AllNetProps
 ) {
     @API("/")
@@ -109,7 +107,7 @@ class AllNet(
         // Proper keychip authentication
         if (props.checkKeychip) {
             // If it's a user keychip, it should be in user database
-            val u = userRepo.findByKeychip(serial)
+            val u = findUserByIncomingSerial(serial)
             if (u != null) {
                 // Create a new session for the user
                 logger.info("> Keychip authenticated: ${u.auId} ${u.computedName}")
@@ -120,11 +118,6 @@ class AllNet(
                 session = keychipSessionService.new(u, reqMap["game_id"] ?: "").token
             }
 
-            // Check if it's a whitelisted keychip
-            else if (!serial.isEmpty() && keychipRepo.existsByKeychipId(serial)) {
-                session = keychipSessionService.new(null, reqMap["game_id"] ?: "").token
-            }
-
             else if (props.keychipPermissiveForTesting) {
                 logger.warn("> Accepted invalid keychip $serial in permissive mode")
                 session = keychipSessionService.new(null, reqMap["game_id"] ?: "").token
@@ -175,6 +168,17 @@ class AllNet(
         return resp.toUrl() + "\n"
     }
 
+    private fun findUserByIncomingSerial(serial: String) = when (serial.length) {
+        FULL_KEYCHIP_LENGTH -> userKeychipRepo.findByKeychipId(serial)?.user
+        SHORT_KEYCHIP_LENGTH -> {
+            // segatools only sends the first 11 characters of the keychip
+            // First, try to find it by suffixing AquaDX's generated suffix, then fall back to matching without the suffixed 4 digits
+            userKeychipRepo.findByKeychipId(serial + KEYCHIP_SUFFIX)?.user
+                ?: userKeychipRepo.findByKeychipIdStartingWith(serial)?.user
+        }
+        else -> userKeychipRepo.findByKeychipId(serial)?.user
+    }
+
     private fun switchUri(hereAddr: Str, localPort: Str, gameId: Str, ver: Str, session: Str?): Str {
         val addr = hereAddr + (if (props.hidePort) "" else ":${props.port ?: localPort}")
 
@@ -198,6 +202,9 @@ class AllNet(
     }
 
     companion object {
+        const val SHORT_KEYCHIP_LENGTH = 11
+        const val FULL_KEYCHIP_LENGTH = 15
+        const val KEYCHIP_SUFFIX = "1337"
         val logger = logger()
     }
 }

src/main/java/icu/samnyan/aqua/sega/allnet/AllNetSecure.kt

@@ -25,7 +25,6 @@ import org.springframework.web.servlet.config.annotation.WebMvcConfigurer
 @Component
 @ConditionalOnBean(AllNetSecureInit::class)
 class TokenChecker(
-    val keyChipRepo: KeyChipRepo,
     val keychipSessionService: KeychipSessionService,
     val frontierProps: FrontierProps,
     val geoip: GeoIP
@@ -56,7 +55,7 @@ class TokenChecker(
         // The token can either be a keychip id (old method) or a session id (new method)
         // Or the frontier token
         val session = keychipSessionService.find(token)
-        if (token.isNotBlank() && (keyChipRepo.existsByKeychipId(token) || session != null
+        if (token.isNotBlank() && (session != null
                 || (frontierProps.enabled && frontierProps.ftk == token)))
         {
             currentSession.set(session)

src/main/java/icu/samnyan/aqua/sega/allnet/Keychip.kt

@@ -1,32 +0,0 @@
-package icu.samnyan.aqua.sega.allnet
-
-import jakarta.persistence.Column
-import jakarta.persistence.Entity
-import jakarta.persistence.Id
-import jakarta.persistence.Table
-import org.springframework.data.jpa.repository.JpaRepository
-import org.springframework.stereotype.Repository
-import java.io.Serializable
-
-/**
- * This is the old method of securing requests - a keychip whitelist,
- * it's kept here only for backwards compatibility.
- */
-@Entity
-@Table(name = "allnet_keychips")
-class Keychip(
-    @Id
-    val id: Long = 0,
-
-    @Column(unique = true, nullable = false)
-    val keychipId: String = "",
-) : Serializable {
-    companion object {
-        const val serialVersionUID = 1L
-    }
-}
-
-@Repository("KeyChipRepository")
-interface KeyChipRepo : JpaRepository<Keychip?, Long?> {
-    fun existsByKeychipId(keychipId: String?): Boolean
-}

src/main/java/icu/samnyan/aqua/sega/allnet/UserKeychip.kt

@@ -0,0 +1,32 @@
+package icu.samnyan.aqua.sega.allnet
+
+import icu.samnyan.aqua.net.db.AquaNetUser
+import jakarta.persistence.*
+import jakarta.transaction.Transactional
+import org.springframework.data.jpa.repository.JpaRepository
+import org.springframework.stereotype.Repository
+
+@Entity
+@Table(name = "user_keychip")
+class UserKeychip(
+    @Id @GeneratedValue(strategy = GenerationType.IDENTITY)
+    val id: Long = 0,
+
+    @ManyToOne
+    @JoinColumn(name = "au_id", nullable = false)
+    var user: AquaNetUser,
+
+    @Column(unique = true, nullable = false, length = 32)
+    val keychipId: String,
+)
+
+@Repository
+interface UserKeychipRepo : JpaRepository<UserKeychip, Long> {
+    fun findByKeychipId(keychipId: String): UserKeychip?
+    fun findByKeychipIdStartingWith(keychipIdPrefix: String): UserKeychip?
+    fun existsByKeychipId(keychipId: String): Boolean
+    fun findAllByUserAuId(auId: Long): List<UserKeychip>
+
+    @Transactional
+    fun deleteByKeychipIdAndUserAuId(keychipId: String, auId: Long): Long
+}