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';
const SAVE_RETRY_DELAYS_MS = [25, 75, 150, 300, 600];
const RETRYABLE_SAVE_ERROR_CODES = new Set(['EPERM', 'EACCES', 'EBUSY']);

let saveQueue: Promise<void> = Promise.resolve();

function wait(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function isRetryableSaveError(error: unknown): boolean {
  if (!error || typeof error !== 'object') {
    return false;
  }

  const code = (error as NodeJS.ErrnoException).code;

  return typeof code === 'string' && RETRYABLE_SAVE_ERROR_CODES.has(code);
}

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 replaceDatabaseFile(tmpPath: string): Promise<void> {
  for (let attempt = 0; ; attempt++) {
    try {
      await fsp.rename(tmpPath, dbFilePath);
      return;
    } catch (error) {
      const delay = SAVE_RETRY_DELAYS_MS[attempt];

      if (!isRetryableSaveError(error) || delay === undefined) {
        throw error;
      }

      await wait(delay);
    }
  }
}

async function writeDatabaseSnapshot(snapshot: Buffer): Promise<void> {
  const tmpPath = dbFilePath + '.tmp-' + randomBytes(6).toString('hex');

  try {
    await fsp.writeFile(tmpPath, snapshot);
    await replaceDatabaseFile(tmpPath);
  } catch (err) {
    await fsp.unlink(tmpPath).catch(() => {});
    throw err;
  }
}

async function atomicSave(data: Uint8Array): Promise<void> {
  const snapshot = Buffer.from(data);
  const saveTask = saveQueue.then(
    () => writeDatabaseSnapshot(snapshot),
    () => writeDatabaseSnapshot(snapshot)
  );

  saveQueue = saveTask.catch(() => {});

  return saveTask;
}

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;
    }
  }
}