From 2b37235666e8998568ff1ad6d40217d4b950fbf9 Mon Sep 17 00:00:00 2001
From: Sketch <75850871+SketchMaster2001@users.noreply.github.com>
Date: Sun, 14 Jan 2024 13:34:09 -0500
Subject: [PATCH] Add a whitelist for testing

---
 ias.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/ias.go b/ias.go
index 88c44d6..36a54c9 100644
--- a/ias.go
+++ b/ias.go
@@ -18,6 +18,7 @@
 package main
 
 import (
+	"bufio"
 	"crypto/md5"
 	"errors"
 	"fmt"
@@ -26,6 +27,8 @@ import (
 	"github.com/jackc/pgx/v4"
 	"log"
 	"math/rand"
+	"os"
+	"slices"
 	"strconv"
 )
 
@@ -34,7 +37,7 @@ const (
 		(device_id, device_token, device_token_hashed, account_id, region, serial_number) 
 	VALUES ($1, $2, $3, $4, $5, $6)`
 	SyncUserStatement = `SELECT 
-		account_id, device_token 
+		account_id, device_token, serial_number
 	FROM userbase WHERE 
 		region = $1 AND
 		device_id = $2`
@@ -92,16 +95,47 @@ func getRegistrationInfo(e *Envelope) {
 	e.AddKVNode("Currency", "POINTS")
 }
 
+func getWhitelistedSerialNumbers() []string {
+	file, err := os.Open("whitelist.txt")
+	if err != nil {
+		panic(err)
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+
+	// Read each line
+	var sns []string
+	for scanner.Scan() {
+		sns = append(sns, scanner.Text())
+
+	}
+
+	// Check for errors during scanning
+	if err = scanner.Err(); err != nil {
+		panic(err)
+	}
+
+	return sns
+}
+
 func syncRegistration(e *Envelope) {
 	var accountId int64
 	var deviceToken string
+	var serialNumber string
 
 	user := pool.QueryRow(ctx, SyncUserStatement, e.Region(), e.DeviceId())
-	err := user.Scan(&accountId, &deviceToken)
+	err := user.Scan(&accountId, &deviceToken, &serialNumber)
 	if err != nil {
 		e.Error(7, "An error occurred querying the database.", err)
 	}
 
+	if !slices.Contains(getWhitelistedSerialNumbers(), serialNumber) {
+		// Since HTTP server runs on a separate Goroutine, this won't shut off the server,
+		// rather kill communication with the requesting console
+		panic(err)
+	}
+
 	e.AddKVNode("AccountId", strconv.FormatInt(accountId, 10))
 	e.AddKVNode("DeviceToken", deviceToken)
 	e.AddKVNode("DeviceTokenExpired", "false")
@@ -133,6 +167,12 @@ func register(e *Envelope) {
 		return
 	}
 
+	if !slices.Contains(getWhitelistedSerialNumbers(), serialNo) {
+		// Since HTTP server runs on a separate Goroutine, this won't shut off the server,
+		// rather kill communication with the requesting console
+		panic(err)
+	}
+
 	// Validate given friend code.
 	userId, err := strconv.ParseUint(deviceCode, 10, 64)
 	if err != nil {