From b00ed4c0b99c27386d845926f6842bdd03a2507e Mon Sep 17 00:00:00 2001
From: maru <maru@myyahoo.com>
Date: Mon, 22 Apr 2024 18:01:17 -0400
Subject: [PATCH] Use WriteFile instead of stream writing to files

---
 api/account/register.go |  6 +++++
 api/savedata/update.go  | 49 ++++++++++++++++++-----------------------
 2 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/api/account/register.go b/api/account/register.go
index ff497c6..c4be2e3 100644
--- a/api/account/register.go
+++ b/api/account/register.go
@@ -3,6 +3,7 @@ package account
 import (
 	"crypto/rand"
 	"fmt"
+	"os"
 
 	"github.com/pagefaultgames/pokerogue-server/db"
 )
@@ -34,5 +35,10 @@ func Register(username, password string) error {
 		return fmt.Errorf("failed to add account record: %s", err)
 	}
 
+	err = os.MkdirAll(fmt.Sprintf("userdata/%x", uuid), 0755)
+	if err != nil && !os.IsExist(err) {
+		return fmt.Errorf(fmt.Sprintf("failed to create userdata folder: %s", err))
+	}
+
 	return nil
 }
diff --git a/api/savedata/update.go b/api/savedata/update.go
index 7b6edc6..a836c28 100644
--- a/api/savedata/update.go
+++ b/api/savedata/update.go
@@ -1,6 +1,7 @@
 package savedata
 
 import (
+	"bytes"
 	"encoding/gob"
 	"fmt"
 	"log"
@@ -19,6 +20,15 @@ func Update(uuid []byte, slot int, save any) error {
 		log.Print("failed to update account last activity")
 	}
 
+	// ideally should have been done at account creation
+	err = os.MkdirAll(fmt.Sprintf("userdata/%x", uuid), 0755)
+	if err != nil && !os.IsExist(err) {
+		return fmt.Errorf(fmt.Sprintf("failed to create userdata folder: %s", err))
+	}
+
+	var filename string
+	var buf bytes.Buffer
+
 	switch save := save.(type) {
 	case defs.SystemSaveData: // System
 		if save.TrainerId == 0 && save.SecretId == 0 {
@@ -33,20 +43,10 @@ func Update(uuid []byte, slot int, save any) error {
 		if err != nil {
 			return fmt.Errorf("failed to update account stats: %s", err)
 		}
+		
+		filename = "system"
 
-		err = os.MkdirAll(fmt.Sprintf("userdata/%x", uuid), 0755)
-		if err != nil && !os.IsExist(err) {
-			return fmt.Errorf("failed to create userdata folder: %s", err)
-		}
-
-		file, err := os.OpenFile(fmt.Sprintf("userdata/%x/system.pzs", uuid), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
-		if err != nil {
-			return fmt.Errorf("failed to open save file for writing: %s", err)
-		}
-
-		defer file.Close()
-
-		zstdEncoder, err := zstd.NewWriter(file)
+		zstdEncoder, err := zstd.NewWriter(&buf)
 		if err != nil {
 			return fmt.Errorf("failed to create zstd encoder: %s", err)
 		}
@@ -64,24 +64,12 @@ func Update(uuid []byte, slot int, save any) error {
 			return fmt.Errorf("slot id %d out of range", slot)
 		}
 
-		fileName := "session"
+		filename = "session"
 		if slot != 0 {
-			fileName += strconv.Itoa(slot)
+			filename += strconv.Itoa(slot)
 		}
 
-		err = os.MkdirAll(fmt.Sprintf("userdata/%x", uuid), 0755)
-		if err != nil && !os.IsExist(err) {
-			return fmt.Errorf(fmt.Sprintf("failed to create userdata folder: %s", err))
-		}
-
-		file, err := os.OpenFile(fmt.Sprintf("userdata/%x/%s.pzs", uuid, fileName), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
-		if err != nil {
-			return fmt.Errorf("failed to open save file for writing: %s", err)
-		}
-
-		defer file.Close()
-
-		zstdEncoder, err := zstd.NewWriter(file)
+		zstdEncoder, err := zstd.NewWriter(&buf)
 		if err != nil {
 			return fmt.Errorf("failed to create zstd encoder: %s", err)
 		}
@@ -96,5 +84,10 @@ func Update(uuid []byte, slot int, save any) error {
 		return fmt.Errorf("invalid data type")
 	}
 
+	err = os.WriteFile(fmt.Sprintf("userdata/%x/%s.pzs", uuid, filename), buf.Bytes(), 0644)
+	if err != nil {
+		return fmt.Errorf("failed to write save to disk: %s", err)
+	}
+
 	return nil
 }