Toju/electron/db/database.ts

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<void> {
  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<void> {
  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<void> {
  if (applicationDataSource?.isInitialized) {
    try {
      await applicationDataSource.destroy();
      console.log('[DB] Connection closed');
    } catch (error) {
      console.error('[DB] Error closing connection:', error);
    } finally {
      applicationDataSource = undefined;
    }
  }
}