import { randomBytes } from 'crypto'; import { app } from 'electron'; import * as fs from 'fs'; import * as fsp from 'fs/promises'; import * as path from 'path'; import { DataSource } from 'typeorm'; import { MessageEntity, UserEntity, RoomEntity, RoomChannelEntity, RoomMemberEntity, RoomRoleEntity, RoomUserRoleEntity, RoomChannelPermissionEntity, ReactionEntity, BanEntity, AttachmentEntity, MetaEntity } from '../entities'; import { settings } from '../settings'; let applicationDataSource: DataSource | undefined; let dbFilePath = ''; let dbBackupPath = ''; // SQLite files start with this 16-byte header string. const SQLITE_MAGIC = 'SQLite format 3\0'; export function getDataSource(): DataSource | undefined { return applicationDataSource; } /** * Returns true when `data` looks like a valid SQLite file * (correct header magic and at least one complete page). */ function isValidSqlite(data: Uint8Array): boolean { if (data.length < 100) return false; const header = Buffer.from(data.buffer, data.byteOffset, 16).toString('ascii'); return header === SQLITE_MAGIC; } /** * Back up the current DB file so there is always a recovery point. * If the main file is corrupted/empty but a valid backup exists, * restore the backup before the app loads the database. */ function safeguardDbFile(): Uint8Array | undefined { if (!fs.existsSync(dbFilePath)) return undefined; const data = new Uint8Array(fs.readFileSync(dbFilePath)); if (isValidSqlite(data)) { fs.copyFileSync(dbFilePath, dbBackupPath); console.log('[DB] Backed up database to', dbBackupPath); return data; } console.warn(`[DB] ${dbFilePath} appears corrupt (${data.length} bytes) - checking backup`); if (fs.existsSync(dbBackupPath)) { const backup = new Uint8Array(fs.readFileSync(dbBackupPath)); if (isValidSqlite(backup)) { fs.copyFileSync(dbBackupPath, dbFilePath); console.warn('[DB] Restored database from backup', dbBackupPath); return backup; } console.error('[DB] Backup is also invalid - starting with a fresh database'); } else { console.error('[DB] No backup available - starting with a fresh database'); } return undefined; } /** * Write the database to disk atomically: write a temp file first, * then rename it over the real file. rename() is atomic on the same * filesystem, so a crash mid-write can never leave a half-written DB. */ async function atomicSave(data: Uint8Array): Promise { const tmpPath = dbFilePath + '.tmp-' + randomBytes(6).toString('hex'); try { await fsp.writeFile(tmpPath, Buffer.from(data)); await fsp.rename(tmpPath, dbFilePath); } catch (err) { await fsp.unlink(tmpPath).catch(() => {}); throw err; } } export async function initializeDatabase(): Promise { const userDataPath = app.getPath('userData'); const dbDir = path.join(userDataPath, 'metoyou'); await fsp.mkdir(dbDir, { recursive: true }); dbFilePath = path.join(dbDir, settings.databaseName); dbBackupPath = dbFilePath + '.bak'; const database = safeguardDbFile(); applicationDataSource = new DataSource({ type: 'sqljs', database, entities: [ MessageEntity, UserEntity, RoomEntity, RoomChannelEntity, RoomMemberEntity, RoomRoleEntity, RoomUserRoleEntity, RoomChannelPermissionEntity, ReactionEntity, BanEntity, AttachmentEntity, MetaEntity ], migrations: [path.join(__dirname, '..', 'migrations', '*.js'), path.join(__dirname, '..', 'migrations', '*.ts')], synchronize: false, logging: false, autoSave: true, autoSaveCallback: atomicSave }); try { await applicationDataSource.initialize(); console.log('[DB] Connection initialised at:', dbFilePath); try { await applicationDataSource.runMigrations(); console.log('[DB] Migrations executed'); } catch (migErr) { console.error('[DB] Migration error:', migErr); } } catch (error) { console.error('[DB] Initialisation error:', error); } } export async function destroyDatabase(): Promise { if (applicationDataSource?.isInitialized) { try { await applicationDataSource.destroy(); console.log('[DB] Connection closed'); } catch (error) { console.error('[DB] Error closing connection:', error); } finally { applicationDataSource = undefined; } } }