Implement banning and custom error messages

database/login.go

@@ -2,19 +2,27 @@ package database
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"time"
 	"wwfc/common"
 	"wwfc/logging"
 
+	"github.com/jackc/pgx/v4"
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/logrusorgru/aurora/v3"
 )
 
-func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsbrcd string, profileId uint32, ngDeviceId uint32) (User, bool) {
+var (
+	ErrDeviceIDMismatch = errors.New("NG device ID mismatch")
+	ErrProfileBannedTOS = errors.New("Profile is banned for violating the Terms of Service")
+)
+
+func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsbrcd string, profileId uint32, ngDeviceId uint32, ipAddress string, ingamesn string) (User, error) {
 	var exists bool
 	err := pool.QueryRow(ctx, DoesUserExist, userId, gsbrcd).Scan(&exists)
 	if err != nil {
-		return User{}, false
+		return User{}, err
 	}
 
 	user := User{
@@ -32,7 +40,7 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 		err := user.CreateUser(pool, ctx)
 		if err != nil {
 			logging.Error("DATABASE", "Error creating user:", aurora.Cyan(userId), aurora.Cyan(gsbrcd), aurora.Cyan(user.ProfileId), "\nerror:", err.Error())
-			return User{}, false
+			return User{}, err
 		}
 
 		logging.Notice("DATABASE", "Created new GPCM user:", aurora.Cyan(userId), aurora.Cyan(gsbrcd), aurora.Cyan(user.ProfileId))
@@ -42,7 +50,7 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 		var lastName *string
 		err := pool.QueryRow(ctx, GetUserProfileID, userId, gsbrcd).Scan(&user.ProfileId, &expectedNgId, &user.Email, &user.UniqueNick, &firstName, &lastName)
 		if err != nil {
-			panic(err)
+			return User{}, err
 		}
 
 		if firstName != nil {
@@ -53,17 +61,17 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 			user.LastName = *lastName
 		}
 
-		if expectedNgId != nil && user.NgDeviceId != 0 {
+		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)))
-				return User{}, false
+				return User{}, ErrDeviceIDMismatch
 			}
 		} else if ngDeviceId != 0 {
 			user.NgDeviceId = ngDeviceId
 			_, err := pool.Exec(ctx, UpdateUserNGDeviceID, user.ProfileId, ngDeviceId)
 			if err != nil {
-				panic(err)
+				return User{}, err
 			}
 		}
 
@@ -86,5 +94,36 @@ func LoginUserToGPCM(pool *pgxpool.Pool, ctx context.Context, userId uint64, gsb
 		})
 	}
 
-	return user, true
+	// Update the user's last IP address and ingamesn
+	_, err = pool.Exec(ctx, UpdateUserLastIPAddress, user.ProfileId, ipAddress, ingamesn)
+	if err != nil {
+		return User{}, err
+	}
+
+	// Find ban from device ID or IP address
+	var banExists bool
+	var banTOS bool
+	var bannedDeviceId uint32
+	timeNow := time.Now()
+	err = pool.QueryRow(ctx, SearchUserBan, user.ProfileId, user.NgDeviceId, ipAddress, timeNow).Scan(&banExists, &banTOS, &bannedDeviceId)
+	if err != nil {
+		if err != pgx.ErrNoRows {
+			return User{}, err
+		}
+
+		banExists = false
+	}
+
+	if banExists {
+		if banTOS {
+			logging.Warn("DATABASE", "Profile", aurora.Cyan(user.ProfileId), "is banned")
+			return User{RestrictedDeviceId: bannedDeviceId}, ErrProfileBannedTOS
+		}
+
+		logging.Warn("DATABASE", "Profile", aurora.Cyan(user.ProfileId), "is restricted")
+		user.Restricted = true
+		user.RestrictedDeviceId = bannedDeviceId
+	}
+
+	return user, nil
 }

database/schema.go

@@ -0,0 +1,22 @@
+package database
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v4/pgxpool"
+)
+
+func UpdateTables(pool *pgxpool.Pool, ctx context.Context) {
+	pool.Exec(ctx, `
+ALTER TABLE ONLY public.users
+	ADD IF NOT EXISTS last_ip_address character varying DEFAULT ''::character varying,
+	ADD IF NOT EXISTS last_ingamesn character varying DEFAULT ''::character varying,
+	ADD IF NOT EXISTS has_ban boolean DEFAULT false,
+	ADD IF NOT EXISTS ban_issued timestamp without time zone,
+	ADD IF NOT EXISTS ban_expires timestamp without time zone,
+	ADD IF NOT EXISTS ban_reason character varying,
+	ADD IF NOT EXISTS ban_reason_hidden character varying,
+	ADD IF NOT EXISTS ban_moderator character varying,
+	ADD IF NOT EXISTS ban_tos boolean
+`)
+}

database/user.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"math/rand"
+	"time"
 
 	"github.com/jackc/pgx/v4/pgxpool"
 )
@@ -19,20 +20,26 @@ const (
 	IsProfileIDInUse        = `SELECT EXISTS(SELECT 1 FROM users WHERE profile_id = $1)`
 	DeleteUserSession       = `DELETE FROM sessions WHERE profile_id = $1`
 	GetUserProfileID        = `SELECT profile_id, ng_device_id, email, unique_nick, firstname, lastname 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`
 
 	GetMKWFriendInfoQuery    = `SELECT mariokartwii_friend_info FROM users WHERE profile_id = $1`
 	UpdateMKWFriendInfoQuery = `UPDATE users SET mariokartwii_friend_info = $2 WHERE profile_id = $1`
 )
 
 type User struct {
-	ProfileId  uint32
-	UserId     uint64
-	GsbrCode   string
-	NgDeviceId uint32
-	Email      string
-	UniqueNick string
-	FirstName  string
-	LastName   string
+	ProfileId          uint32
+	UserId             uint64
+	GsbrCode           string
+	NgDeviceId         uint32
+	Email              string
+	UniqueNick         string
+	FirstName          string
+	LastName           string
+	Restricted         bool
+	RestrictedDeviceId uint32
 }
 
 var (
@@ -130,6 +137,24 @@ func GetProfile(pool *pgxpool.Pool, ctx context.Context, profileId uint32) (User
 	return user, true
 }
 
+func BanUser(pool *pgxpool.Pool, ctx context.Context, profileId uint32, tos bool, length time.Duration, reason string, reasonHidden string, moderator string) bool {
+	_, err := pool.Exec(ctx, UpdateUserBan, profileId, time.Now(), time.Now().Add(length), reason, reasonHidden, moderator, tos)
+	if err != nil {
+		return false
+	}
+
+	return true
+}
+
+func UnbanUser(pool *pgxpool.Pool, ctx context.Context, profileId uint32) bool {
+	_, err := pool.Exec(ctx, DisableUserBan, profileId)
+	if err != nil {
+		return false
+	}
+
+	return true
+}
+
 func GetMKWFriendInfo(pool *pgxpool.Pool, ctx context.Context, profileId uint32) string {
 	var info string
 	err := pool.QueryRow(ctx, GetMKWFriendInfoQuery, profileId).Scan(&info)