fix: preserve cold storage during sync and backup

src/ts/characters.ts

@@ -875,7 +875,7 @@ export async function changeChar(index: number, arg:{
     reseter();
     if(DBState.db.characters?.[index]?.coldstorage){
         const coldData = await getColdStorageItem(DBState.db.characters[index].coldstorage!)
-        if(coldData.character && coldData.character.chaId === DBState.db.characters[index].chaId){
+        if(coldData?.character && coldData.character.chaId === DBState.db.characters[index].chaId){
             DBState.db.characters[index] = coldData.character
         }
         else{

src/ts/drive/backuplocal.ts

@@ -9,7 +9,7 @@ import { relaunch } from "@tauri-apps/plugin-process";
 import { decryptBuffer, encryptBuffer, sleep } from "../util";
 import { hubURL } from "../characterCards";
 import { language } from "src/lang";
-import { getColdStorageItem, listColdDataKeys, setColdStorageItem } from "../process/coldstorage.svelte";
+import { collectColdStorageBackupPayloads, getColdStorageBackupKey, getColdStorageItem, isColdStorageBackupData, listColdDataKeys, setColdStorageItem } from "../process/coldstorage.svelte";
 import { DBState } from "../stores.svelte";
 
 function getBasename(data:string){
@@ -19,31 +19,22 @@ function getBasename(data:string){
     return lasts
 }
 
-function getColdStorageBackupKey(name: string): string | null {
-    const match = name.match(/^(?:coldstorage[/_])?([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12})\.json$/)
-    return match?.[1] ?? null
-}
-
-function isColdStorageBackupData(data: unknown): boolean {
-    if (Array.isArray(data)) {
-        return true
-    }
-
-    return !!data
-        && typeof data === 'object'
-        && ('character' in data || 'message' in data)
-}
-
 export async function SaveLocalBackup(){
     alertWait("Saving local backup...")
+    const db = getDatabase()
+    const coldStoragePayloads = await collectColdStorageBackupPayloads(db)
+    if(coldStoragePayloads.missingKeys.length > 0){
+        alertError(`Backup failed. ${coldStoragePayloads.missingKeys.length} cold storage item(s) are missing.`)
+        return
+    }
+
     const writer = new LocalWriter()
     const r = await writer.init()
     if(!r){
         alertError('Failed')
         return
     }
 
-    const db = getDatabase()
     const assetMap = new Map<string, { charName: string, assetName: string }>()
     if (db.characters) {
         for (const char of db.characters) {
@@ -160,21 +151,11 @@ export async function SaveLocalBackup(){
         }
     }
 
-    if(!forageStorage.isAccount){
-        //save coldstorages
-        const coldKeys = await listColdDataKeys()
-        for(let i=0;i<coldKeys.length;i++){
-            const key = coldKeys[i]
-            let message = `Saving local Backup Cold data... (${i + 1} / ${coldKeys.length})`
-            alertWait(message)
-            const data = await getColdStorageItem(key)
-            if(data){
-                const encoded = new TextEncoder().encode(JSON.stringify(data))
-                await writer.writeBackup(`coldstorage_${key}.json`, encoded)
-            } else {
-                missingAssets.push(`coldstorage_${key}.json`)
-            }
-        }
+    for(let i=0;i<coldStoragePayloads.payloads.length;i++){
+        const payload = coldStoragePayloads.payloads[i]
+        let message = `Saving local Backup Cold data... (${i + 1} / ${coldStoragePayloads.payloads.length})`
+        alertWait(message)
+        await writer.writeBackup(payload.backupName, payload.encoded)
     }
 
     const dbWithoutAccount = { ...db, account: undefined }
@@ -235,14 +216,20 @@ export async function SavePartialLocalBackup(){
     }
 
     alertWait("Saving partial local backup...")
+    const db = getDatabase()
+    const coldStoragePayloads = await collectColdStorageBackupPayloads(db)
+    if(coldStoragePayloads.missingKeys.length > 0){
+        alertError(`Partial backup failed. ${coldStoragePayloads.missingKeys.length} cold storage item(s) are missing.`)
+        return
+    }
+
     const writer = new LocalWriter()
     const r = await writer.init()
     if(!r){
         alertError('Failed')
         return
     }
 
-    const db = getDatabase()
     const assetMap = new Map<string, { charName: string, assetName: string }>()
 
     // Only collect main profile images for both characters and groups
@@ -384,21 +371,11 @@ export async function SavePartialLocalBackup(){
         }
     }
 
-    if(!forageStorage.isAccount){
-        //save coldstorages
-        const coldKeys = await listColdDataKeys()
-        for(let i=0;i<coldKeys.length;i++){
-            const key = coldKeys[i]
-            let message = `Saving partial local Backup Cold data... (${i + 1} / ${coldKeys.length})`
-            alertWait(message)
-            const data = await getColdStorageItem(key)
-            if(data){
-                const encoded = new TextEncoder().encode(JSON.stringify(data))
-                await writer.writeBackup(`coldstorage_${key}.json`, encoded)
-            } else {
-                missingAssets.push(`coldstorage_${key}.json`)
-            }
-        }
+    for(let i=0;i<coldStoragePayloads.payloads.length;i++){
+        const payload = coldStoragePayloads.payloads[i]
+        let message = `Saving partial local Backup Cold data... (${i + 1} / ${coldStoragePayloads.payloads.length})`
+        alertWait(message)
+        await writer.writeBackup(payload.backupName, payload.encoded)
     }
 
     const dbWithoutAccount = { ...db, account: undefined }
@@ -448,6 +425,8 @@ export function LoadLocalBackup(){
             const CHUNK_SIZE = 1024 * 1024; // 1MB chunk size
             let bytesRead = 0;
             let remainingBuffer = new Uint8Array();
+            let pendingDatabase: Uint8Array | null = null;
+            const restoredColdStorageKeys = new Set<string>();
 
             while (true) {
                 const { done, value } = await reader.read();
@@ -500,53 +479,25 @@ export function LoadLocalBackup(){
                     }
 
                     else if (name === 'database.risudat') {
-                        let db = new Uint8Array(data);
-                        if(encryptionMeta.type === 'account' && encryptionMeta.time){
-                            try {
-                                const key = (await (await fetch(`https://sv.risuai.xyz/cryptokey?key=${encryptionMeta.time}`)).json()).key
-                                const decrypted = await decryptBuffer(db, key)
-                                db = new Uint8Array(decrypted)
-                            }
-                            catch (e) {
-                                console.error('Failed to decrypt database backup:', e)
-                                alertError('Failed to decrypt database backup, will attempt to load it without decryption.')
-                            }
-                        }
-                        const dbData = await decodeRisuSave(db);
-                        setDatabaseLite(dbData);
-                        requiresFullEncoderReload.state = true;
-                        if (isTauri) {
-                            await writeFile('database/database.bin', db, { baseDir: BaseDirectory.AppData });
-                            await relaunch();
-                            alertStore.set({
-                                type: "wait",
-                                msg: "Success, Refreshing your app."
-                            });
-                        } else {
-                            await forageStorage.setItem('database/database.bin', db);
-                            location.search = '';
-                            alertStore.set({
-                                type: "wait",
-                                msg: "Success, Refreshing your app."
-                            });
-                        }
+                        pendingDatabase = new Uint8Array(data);
                     }
 
                     else {
                         const coldStorageKey = getColdStorageBackupKey(name)
                         let handledAsColdStorage = false
 
-                        if (coldStorageKey && forageStorage.isAccount) {
+                        if (coldStorageKey) {
                             handledAsColdStorage = true
-                        }
-                        else if (coldStorageKey) {
                             try {
                                 const text = new TextDecoder().decode(data)
                                 const jsonData = JSON.parse(text)
 
                                 if (isColdStorageBackupData(jsonData)) {
-                                    await setColdStorageItem(coldStorageKey, jsonData)
-                                    handledAsColdStorage = true
+                                    if(await setColdStorageItem(coldStorageKey, jsonData)){
+                                        restoredColdStorageKeys.add(coldStorageKey)
+                                    } else {
+                                        console.error(`Failed to restore cold storage item ${coldStorageKey}`)
+                                    }
                                 } else {
                                     console.warn(`Skipping invalid cold storage backup item ${name}`)
                                 }
@@ -573,6 +524,57 @@ export function LoadLocalBackup(){
                 remainingBuffer = remainingBuffer.slice(offset);
             }
 
+            if(!pendingDatabase){
+                alertError('Failed, Is file corrupted?')
+                return
+            }
+
+            let db = pendingDatabase;
+            if(encryptionMeta.type === 'account' && encryptionMeta.time){
+                try {
+                    const key = (await (await fetch(`https://sv.risuai.xyz/cryptokey?key=${encryptionMeta.time}`)).json()).key
+                    const decrypted = await decryptBuffer(db, key)
+                    db = new Uint8Array(decrypted)
+                }
+                catch (e) {
+                    console.error('Failed to decrypt database backup:', e)
+                    alertError('Failed to decrypt database backup, will attempt to load it without decryption.')
+                }
+            }
+            const dbData = await decodeRisuSave(db);
+            const missingColdStorageKeys:string[] = []
+            for(const key of await listColdDataKeys(dbData)){
+                if(restoredColdStorageKeys.has(key)){
+                    continue
+                }
+                const existingColdStorage = await getColdStorageItem(key)
+                if(!existingColdStorage){
+                    missingColdStorageKeys.push(key)
+                }
+            }
+            if(missingColdStorageKeys.length > 0){
+                alertError(`Backup restore failed. ${missingColdStorageKeys.length} cold storage item(s) are missing.`)
+                return
+            }
+
+            setDatabaseLite(dbData);
+            requiresFullEncoderReload.state = true;
+            if (isTauri) {
+                await writeFile('database/database.bin', db, { baseDir: BaseDirectory.AppData });
+                await relaunch();
+                alertStore.set({
+                    type: "wait",
+                    msg: "Success, Refreshing your app."
+                });
+            } else {
+                await forageStorage.setItem('database/database.bin', db);
+                location.search = '';
+                alertStore.set({
+                    type: "wait",
+                    msg: "Success, Refreshing your app."
+                });
+            }
+
             alertNormal('Success');
         };
 
@@ -581,4 +583,4 @@ export function LoadLocalBackup(){
         console.error(error);
         alertError('Failed, Is file corrupted?')
     }
-}
+}

src/ts/drive/drive.ts

@@ -8,6 +8,7 @@ import { relaunch } from '@tauri-apps/plugin-process';
 import { sleep } from "../util";
 import { hubURL } from "../characterCards";
 import { decodeRisuSave, encodeRisuSaveLegacy } from "../storage/risuSave";
+import { collectColdStorageBackupPayloads, getColdStorageBackupName, isColdStorageBackupData, listColdDataKeys, setColdStorageItem } from "../process/coldstorage.svelte";
 
 export async function checkDriver(type:'save'|'load'|'loadtauri'|'savetauri'|'reftoken'){
     const CLIENT_ID = '580075990041-l26k2d3c0nemmqiu3d3aag01npfrkn76.apps.googleusercontent.com';
@@ -126,6 +127,12 @@ async function backupDrive(ACCESS_TOKEN:string) {
         return d.name
     })
 
+    const coldStoragePayloads = await collectColdStorageBackupPayloads(getDatabase())
+    if(coldStoragePayloads.missingKeys.length > 0){
+        alertError(`Backup failed. ${coldStoragePayloads.missingKeys.length} cold storage item(s) are missing.`)
+        return
+    }
+
     if(isTauri){
         const assets = await readDir('assets', {baseDir: BaseDirectory.AppData})
         let i = 0;
@@ -164,6 +171,18 @@ async function backupDrive(ACCESS_TOKEN:string) {
         }
     }
 
+    for(let i=0;i<coldStoragePayloads.payloads.length;i++){
+        const payload = coldStoragePayloads.payloads[i]
+        alertStore.set({
+            type: "wait",
+            msg: `Uploading Cold Storage... (${i + 1} / ${coldStoragePayloads.payloads.length})`
+        })
+        if(fileNames.includes(payload.backupName)){
+            continue
+        }
+        await createFileInFolder(ACCESS_TOKEN, payload.backupName, payload.encoded)
+    }
+
     const dbData = encodeRisuSaveLegacy(getDatabase(), 'compression')
 
     alertStore.set({
@@ -280,8 +299,11 @@ async function loadDrive(ACCESS_TOKEN:string, mode: 'backup'|'sync'):Promise<voi
         }
 
         const db:Database = mode === 'backup' ? await getDbFromList() : JSON.parse(Buffer.from(await getFileData(ACCESS_TOKEN, dbs[0][0].id)).toString('utf-8'))
-        lastSaved = Date.now()
-        localStorage.setItem('risu_lastsaved', `${lastSaved}`)
+        const coldStorageRestoreFailures = await restoreColdStorageFromDrive(ACCESS_TOKEN, files, db, mode)
+        if(coldStorageRestoreFailures.length > 0){
+            alertError(`${mode === 'sync' ? 'Sync' : 'Backup restore'} failed. ${coldStorageRestoreFailures.length} cold storage item(s) could not be restored.`)
+            return
+        }
         const requiredImages = (await getUncleanables(db))
         let ind = 0;
         let errorLogs:string[] = []
@@ -337,6 +359,8 @@ async function loadDrive(ACCESS_TOKEN:string, mode: 'backup'|'sync'):Promise<voi
 
         if(isTauri){
             await writeFile('database/database.bin', dbData, {baseDir: BaseDirectory.AppData})
+            lastSaved = Date.now()
+            localStorage.setItem('risu_lastsaved', `${lastSaved}`)
             relaunch()
             alertStore.set({
                 type: "wait",
@@ -345,6 +369,8 @@ async function loadDrive(ACCESS_TOKEN:string, mode: 'backup'|'sync'):Promise<voi
         }
         else{
             await forageStorage.setItem('database/database.bin', dbData)
+            lastSaved = Date.now()
+            localStorage.setItem('risu_lastsaved', `${lastSaved}`)
             location.search = ''
             alertStore.set({
                 type: "wait",
@@ -357,6 +383,50 @@ async function loadDrive(ACCESS_TOKEN:string, mode: 'backup'|'sync'):Promise<voi
     }
 }
 
+async function restoreColdStorageFromDrive(
+    ACCESS_TOKEN:string,
+    files:DriveFile[],
+    db:Database,
+    mode:'backup'|'sync'
+) {
+    const coldKeys = await listColdDataKeys(db)
+    const failures:string[] = []
+    for(let i=0;i<coldKeys.length;i++){
+        const key = coldKeys[i]
+        const names = new Set([
+            getColdStorageBackupName(key),
+            `coldstorage/${key}.json`,
+            `${key}.json`
+        ])
+        const file = files.find((driveFile) => names.has(driveFile.name))
+        if(!file){
+            console.warn(`Cold storage data not found in Drive backup: ${key}`)
+            failures.push(key)
+            continue
+        }
+        alertStore.set({
+            type: "wait",
+            msg: `${mode === 'sync' ? 'Sync' : 'Loading Backup'} Cold Storage... (${i + 1} / ${coldKeys.length})`
+        })
+        try {
+            const jsonData = JSON.parse(new TextDecoder().decode(await getFileData(ACCESS_TOKEN, file.id)))
+            if(isColdStorageBackupData(jsonData)){
+                if(!await setColdStorageItem(key, jsonData)){
+                    failures.push(key)
+                }
+            }
+            else{
+                console.warn(`Skipping invalid cold storage Drive item ${file.name}`)
+                failures.push(key)
+            }
+        } catch (error) {
+            console.error(`Failed to restore cold storage item ${key}:`, error)
+            failures.push(key)
+        }
+    }
+    return failures
+}
+
 function checkImageExist(image:string){
 
 }
@@ -454,4 +524,4 @@ async function getFileData(ACCESS_TOKEN:string,fileId:string) {
     } else {
         throw "Error in response when reading files in folder"
     }
-  }
+  }