From 6176e43fb70cabd6e34c2875b55bd0b5bec8e9e7 Mon Sep 17 00:00:00 2001
From: mkwcat <cat@mkw.re>
Date: Tue, 7 May 2024 05:15:36 -0400
Subject: [PATCH] GameStats: Save and reload state

---
 gamestats/auth.go |  5 +++--
 gamestats/main.go | 48 +++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/gamestats/auth.go b/gamestats/auth.go
index ad7747b..3828df5 100644
--- a/gamestats/auth.go
+++ b/gamestats/auth.go
@@ -21,7 +21,8 @@ func (g *GameStatsSession) auth(command common.GameSpyCommand) {
 
 	// TODO: Validate "response"
 	g.SessionKey = rand.Int31n(290000000) + 10000000
-	g.GameInfo = game
+	g.GameName = command.OtherValues["gamename"]
+	g.gameInfo = game
 
 	g.Write(common.GameSpyCommand{
 		Command:      "lc",
@@ -86,7 +87,7 @@ func (g *GameStatsSession) authp(command common.GameSpyCommand) {
 	g.ModuleName = "GSTATS:" + strconv.FormatInt(int64(g.User.ProfileId), 10)
 	g.Authenticated = true
 
-	logging.Notice(g.ModuleName, "Authenticated, game name:", aurora.Cyan(g.GameInfo.Name))
+	logging.Notice(g.ModuleName, "Authenticated, game name:", aurora.Cyan(g.gameInfo.Name))
 
 	g.Write(common.GameSpyCommand{
 		Command:      "pauthr",
diff --git a/gamestats/main.go b/gamestats/main.go
index a20226d..b0b45b9 100644
--- a/gamestats/main.go
+++ b/gamestats/main.go
@@ -2,7 +2,9 @@ package gamestats
 
 import (
 	"context"
+	"encoding/gob"
 	"fmt"
+	"os"
 	"strings"
 	"wwfc/common"
 	"wwfc/database"
@@ -23,7 +25,8 @@ type GameStatsSession struct {
 	Challenge  string
 
 	SessionKey int32
-	GameInfo   *common.GameInfo
+	GameName   string
+	gameInfo   *common.GameInfo
 
 	Authenticated bool
 	LoginID       int
@@ -64,9 +67,50 @@ func StartServer(reload bool) {
 	if err != nil {
 		panic(err)
 	}
+
+	if reload {
+		// Load state
+		file, err := os.Open("/state/gstats_sessions.gob")
+		if err != nil {
+			panic(err)
+		}
+
+		decoder := gob.NewDecoder(file)
+
+		err = decoder.Decode(&sessionsByConnIndex)
+		file.Close()
+
+		if err != nil {
+			panic(err)
+		}
+
+		for _, session := range sessionsByConnIndex {
+			session.gameInfo = common.GetGameInfoByName(session.GameName)
+			if session.gameInfo == nil {
+				logging.Error(session.ModuleName, "Unknown game from reload:", aurora.Cyan(session.GameName))
+				// Force close the session now to prevent a panic later
+				common.CloseConnection(ServerName, session.ConnIndex)
+				delete(sessionsByConnIndex, session.ConnIndex)
+			}
+		}
+	}
 }
 
 func Shutdown() {
+	// Save state
+	file, err := os.OpenFile("/state/gstats_sessions.gob", os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		panic(err)
+	}
+
+	encoder := gob.NewEncoder(file)
+
+	err = encoder.Encode(sessionsByConnIndex)
+	file.Close()
+
+	if err != nil {
+		panic(err)
+	}
 }
 
 func NewConnection(index uint64, address string) {
@@ -77,7 +121,7 @@ func NewConnection(index uint64, address string) {
 		Challenge:  common.RandomString(10),
 
 		SessionKey: 0,
-		GameInfo:   nil,
+		gameInfo:   nil,
 
 		Authenticated: false,
 		LoginID:       0,