From 2e798b591a944881b609dccd0244e65cb4eb7ed0 Mon Sep 17 00:00:00 2001
From: Palapeli <26661008+mkwcat@users.noreply.github.com>
Date: Sat, 1 Mar 2025 03:49:59 -0500
Subject: [PATCH] Database: Support multiple devices on one profile

---
 common/config.go  | 14 +++++++--
 database/login.go | 77 ++++++++++++++++++++++++++++++++++++++---------
 database/user.go  | 33 +++++++++++++-------
 gpcm/error.go     |  4 +--
 gpcm/login.go     |  2 ++
 schema.sql        | 12 +++++++-
 6 files changed, 111 insertions(+), 31 deletions(-)

diff --git a/common/config.go b/common/config.go
index 66ea07f..6960501 100644
--- a/common/config.go
+++ b/common/config.go
@@ -40,19 +40,29 @@ type Config struct {
 
 	APISecret string `xml:"apiSecret"`
 
-	AllowDefaultDolphinKeys bool `xml:"allowDefaultDolphinKeys"`
+	AllowDefaultDolphinKeys     bool `xml:"allowDefaultDolphinKeys"`
+	AllowMultipleDeviceIDs      bool `xml:"allowMultipleDeviceIDs"`
+	AllowConnectWithoutDeviceID bool `xml:"allowConnectWithoutDeviceID"`
 
 	ServerName string `xml:"serverName,omitempty"`
 }
 
+var config Config
+var configLoaded bool
+
 func GetConfig() Config {
+	if configLoaded {
+		return config
+	}
+
 	data, err := os.ReadFile("config.xml")
 	if err != nil {
 		panic(err)
 	}
 
-	var config Config
 	config.AllowDefaultDolphinKeys = true
+	config.AllowMultipleDeviceIDs = false
+	config.AllowConnectWithoutDeviceID = false
 	config.ServerName = "WiiLink"
 
 	err = xml.Unmarshal(data, &config)
diff --git a/database/login.go b/database/login.go
index c9e3d87..511bef4 100644
--- a/database/login.go
+++ b/database/login.go
@@ -32,7 +32,10 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 
 	if !exists {
 		user.ProfileId = profileId
-		user.NgDeviceId = ngDeviceId
+		user.NgDeviceId = []uint32{ngDeviceId}
+		if ngDeviceId == 0 {
+			user.NgDeviceId = []uint32{}
+		}
 		user.UniqueNick = common.Base32Encode(userId) + gsbrcd
 		user.Email = user.UniqueNick + "@nds"
 
@@ -45,10 +48,9 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 
 		logging.Notice("DATABASE", "Created new GPCM user:", aurora.Cyan(userId), aurora.Cyan(gsbrcd), aurora.Cyan(user.ProfileId))
 	} else {
-		var expectedNgId *uint32
 		var firstName *string
 		var lastName *string
-		err := pool.QueryRow(ctx, GetUserProfileID, userId, gsbrcd).Scan(&user.ProfileId, &expectedNgId, &user.Email, &user.UniqueNick, &firstName, &lastName, &user.OpenHost)
+		err := pool.QueryRow(ctx, GetUserProfileID, userId, gsbrcd).Scan(&user.ProfileId, &user.NgDeviceId, &user.Email, &user.UniqueNick, &firstName, &lastName, &user.OpenHost)
 		if err != nil {
 			return User{}, err
 		}
@@ -61,18 +63,44 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 			user.LastName = *lastName
 		}
 
-		if expectedNgId != nil && *expectedNgId != 0 {
-			user.NgDeviceId = *expectedNgId
-			if ngDeviceId != 0 && user.NgDeviceId != ngDeviceId {
-				logging.Error("DATABASE", "NG device ID mismatch for profile", aurora.Cyan(user.ProfileId), "- expected", aurora.Cyan(fmt.Sprintf("%08x", user.NgDeviceId)), "but got", aurora.Cyan(fmt.Sprintf("%08x", ngDeviceId)))
+		validDeviceId := false
+		deviceIdList := ""
+		for index, id := range user.NgDeviceId {
+			if id == ngDeviceId {
+				validDeviceId = true
+				break
+			}
+
+			if id == 0 {
+				// Replace the 0 with the actual device ID
+				user.NgDeviceId[index] = ngDeviceId
+				_, err = pool.Exec(ctx, UpdateUserNGDeviceID, user.ProfileId, user.NgDeviceId)
+				validDeviceId = true
+				break
+			}
+
+			deviceIdList += aurora.Cyan(fmt.Sprintf("%08x", id)).String() + ", "
+		}
+
+		if !validDeviceId && ngDeviceId != 0 {
+			if len(user.NgDeviceId) > 0 && !common.GetConfig().AllowMultipleDeviceIDs {
+				logging.Error("DATABASE", "NG device ID mismatch for profile", aurora.Cyan(user.ProfileId), "- expected one of {", deviceIdList[:len(deviceIdList)-2], "} but got", aurora.Cyan(fmt.Sprintf("%08x", ngDeviceId)))
+				return User{}, ErrDeviceIDMismatch
+			} else if len(user.NgDeviceId) > 0 {
+				logging.Warn("DATABASE", "Adding NG device ID", aurora.Cyan(fmt.Sprintf("%08x", ngDeviceId)), "to profile", aurora.Cyan(user.ProfileId))
+			}
+
+			user.NgDeviceId = append(user.NgDeviceId, ngDeviceId)
+			_, err = pool.Exec(ctx, UpdateUserNGDeviceID, user.ProfileId, user.NgDeviceId)
+		} else if !validDeviceId && ngDeviceId == 0 {
+			if len(user.NgDeviceId) > 0 && !common.GetConfig().AllowConnectWithoutDeviceID {
+				logging.Error("DATABASE", "NG device ID not provided for profile", aurora.Cyan(user.ProfileId), "- expected one of {", deviceIdList[:len(deviceIdList)-2], "} but got", aurora.Cyan("00000000"))
 				return User{}, ErrDeviceIDMismatch
 			}
-		} else if ngDeviceId != 0 {
-			user.NgDeviceId = ngDeviceId
-			_, err := pool.Exec(ctx, UpdateUserNGDeviceID, user.ProfileId, ngDeviceId)
-			if err != nil {
-				return User{}, err
-			}
+		}
+
+		if err != nil {
+			return User{}, err
 		}
 
 		if profileId != 0 && user.ProfileId != profileId {
@@ -103,9 +131,9 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 	// Find ban from device ID or IP address
 	var banExists bool
 	var banTOS bool
-	var bannedDeviceId uint32
+	var bannedDeviceIdList []uint32
 	timeNow := time.Now()
-	err = pool.QueryRow(ctx, SearchUserBan, user.ProfileId, user.NgDeviceId, ipAddress, timeNow).Scan(&banExists, &banTOS, &bannedDeviceId)
+	err = pool.QueryRow(ctx, SearchUserBan, user.NgDeviceId, user.ProfileId, ipAddress, timeNow).Scan(&banExists, &banTOS, &bannedDeviceIdList)
 	if err != nil {
 		if err != pgx.ErrNoRows {
 			return User{}, err
@@ -115,6 +143,25 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 	}
 
 	if banExists {
+		// Find first device ID in common
+		bannedDeviceId := uint32(0)
+		for _, id := range bannedDeviceIdList {
+			for _, id2 := range user.NgDeviceId {
+				if id == id2 {
+					bannedDeviceId = id
+					break
+				}
+			}
+
+			if bannedDeviceId != 0 {
+				break
+			}
+		}
+
+		if bannedDeviceId == 0 && len(bannedDeviceIdList) > 0 {
+			bannedDeviceId = bannedDeviceIdList[len(bannedDeviceIdList)-1]
+		}
+
 		if banTOS {
 			logging.Warn("DATABASE", "Profile", aurora.Cyan(user.ProfileId), "is banned")
 			return User{RestrictedDeviceId: bannedDeviceId}, ErrProfileBannedTOS
diff --git a/database/user.go b/database/user.go
index 335202c..ce6f6df 100644
--- a/database/user.go
+++ b/database/user.go
@@ -22,9 +22,29 @@ const (
 	GetUserProfileID        = `SELECT profile_id, ng_device_id, email, unique_nick, firstname, lastname, open_host FROM users WHERE user_id = $1 AND gsbrcd = $2`
 	UpdateUserLastIPAddress = `UPDATE users SET last_ip_address = $2, last_ingamesn = $3 WHERE profile_id = $1`
 	UpdateUserBan           = `UPDATE users SET has_ban = true, ban_issued = $2, ban_expires = $3, ban_reason = $4, ban_reason_hidden = $5, ban_moderator = $6, ban_tos = $7 WHERE profile_id = $1`
-	SearchUserBan           = `SELECT has_ban, ban_tos, ng_device_id FROM users WHERE has_ban = true AND (profile_id = $1 OR ng_device_id = $2 OR last_ip_address = $3) AND (ban_expires IS NULL OR ban_expires > $4) ORDER BY ban_tos DESC LIMIT 1`
 	DisableUserBan          = `UPDATE users SET has_ban = false WHERE profile_id = $1`
 
+	SearchUserBan = `WITH known_ng_device_ids AS (
+		WITH RECURSIVE device_tree AS (
+			SELECT unnest(ng_device_id) AS device_id
+				FROM users
+				WHERE ng_device_id && $1
+			UNION
+			SELECT unnest(ng_device_id)
+				FROM users
+				JOIN device_tree dt
+				ON ng_device_id && array[dt.device_id]
+		) SELECT array_agg(DISTINCT device_id) FROM device_tree
+	)
+	SELECT has_ban, ban_tos, ng_device_id 
+		FROM users
+		WHERE has_ban = true
+			AND (profile_id = $2
+				OR ng_device_id && (SELECT * FROM known_ng_device_ids)
+				OR last_ip_address = $3)
+			AND (ban_expires IS NULL OR ban_expires > $4)
+			ORDER BY ban_tos DESC LIMIT 1`
+
 	GetMKWFriendInfoQuery    = `SELECT mariokartwii_friend_info FROM users WHERE profile_id = $1`
 	UpdateMKWFriendInfoQuery = `UPDATE users SET mariokartwii_friend_info = $2 WHERE profile_id = $1`
 )
@@ -33,7 +53,7 @@ type User struct {
 	ProfileId          uint32
 	UserId             uint64
 	GsbrCode           string
-	NgDeviceId         uint32
+	NgDeviceId         []uint32
 	Email              string
 	UniqueNick         string
 	FirstName          string
@@ -94,15 +114,6 @@ func (user *User) UpdateProfileID(pool *pgxpool.Pool, ctx context.Context, newPr
 	return err
 }
 
-func (user *User) UpdateDeviceID(pool *pgxpool.Pool, ctx context.Context, newDeviceId uint32) error {
-	_, err := pool.Exec(ctx, UpdateUserNGDeviceID, user.ProfileId, newDeviceId)
-	if err == nil {
-		user.NgDeviceId = newDeviceId
-	}
-
-	return err
-}
-
 func GetUniqueUserID() uint64 {
 	// Not guaranteed unique but doesn't matter in practice if multiple people have the same user ID.
 	return uint64(rand.Int63n(0x80000000000))
diff --git a/gpcm/error.go b/gpcm/error.go
index 84323a1..328f31c 100644
--- a/gpcm/error.go
+++ b/gpcm/error.go
@@ -430,8 +430,8 @@ func (g *GameSpySession) replyError(err GPError) {
 	}
 
 	deviceId := g.User.RestrictedDeviceId
-	if deviceId == 0 {
-		deviceId = g.User.NgDeviceId
+	if deviceId == 0 && len(g.User.NgDeviceId) > 0 {
+		deviceId = g.User.NgDeviceId[0]
 	}
 
 	msg := err.GetMessageTranslate(g.GameName, g.Region, g.Language, g.ConsoleFriendCode, deviceId)
diff --git a/gpcm/login.go b/gpcm/login.go
index 57072cf..d6c8d2e 100644
--- a/gpcm/login.go
+++ b/gpcm/login.go
@@ -439,6 +439,8 @@ func (g *GameSpySession) performLoginWithDatabase(userId uint64, gsbrCode string
 	g.User = user
 
 	if err != nil {
+		logging.Error(g.ModuleName, "DB error:", err)
+
 		if err == database.ErrProfileIDInUse {
 			g.replyError(GPError{
 				ErrorCode:   ErrLogin.ErrorCode,
diff --git a/schema.sql b/schema.sql
index 0796d88..79fa1c2 100644
--- a/schema.sql
+++ b/schema.sql
@@ -50,6 +50,16 @@ ALTER TABLE ONLY public.users
     ADD IF NOT EXISTS ban_tos boolean,
 	ADD IF NOT EXISTS open_host boolean DEFAULT false;
 
+--
+-- Change ng_device_id from bigint to bigint[]
+--
+DO $$ 
+BEGIN
+    IF (SELECT data_type FROM information_schema.columns WHERE table_name='users' AND column_name='ng_device_id') != 'ARRAY' THEN
+        ALTER TABLE public.users
+            ALTER COLUMN ng_device_id TYPE bigint[] using array[ng_device_id];
+    END IF;
+END $$;
 
 ALTER TABLE public.users OWNER TO wiilink;
 
@@ -82,7 +92,7 @@ ALTER TABLE public.mario_kart_wii_sake OWNER TO wiilink;
 
 CREATE SEQUENCE IF NOT EXISTS public.users_profile_id_seq
     AS integer
-    START WITH 1
+    START WITH 1000000000
     INCREMENT BY 1
     NO MINVALUE
     NO MAXVALUE