Toju/server/src/db/database.ts

import { randomBytes } from 'crypto';
import fs from 'fs';
import fsp from 'fs/promises';
import path from 'path';
import { DataSource } from 'typeorm';
import {
  AuthUserEntity,
  ServerEntity,
  ServerTagEntity,
  ServerChannelEntity,
  ServerRoleEntity,
  ServerUserRoleEntity,
  ServerChannelPermissionEntity,
  JoinRequestEntity,
  ServerMembershipEntity,
  ServerInviteEntity,
  ServerBanEntity,
  GameMatchMissEntity,
  ServerPluginRequirementEntity,
  ServerPluginEventDefinitionEntity,
  PluginDataEntity,
  ServerPluginSettingsEntity,
  PluginUserMetadataEntity
} from '../entities';
import { serverMigrations } from '../migrations';
import {
  findExistingPath,
  isPackagedRuntime,
  resolvePersistentDataPath,
  resolveRuntimePath
} from '../runtime-paths';

const LEGACY_PACKAGED_DB_FILE = path.join(resolveRuntimePath('data'), 'metoyou.sqlite');
const LEGACY_PACKAGED_DB_BACKUP = LEGACY_PACKAGED_DB_FILE + '.bak';

function resolveDefaultDbFile(): string {
  return isPackagedRuntime()
    ? resolvePersistentDataPath('metoyou.sqlite')
    : LEGACY_PACKAGED_DB_FILE;
}

function resolveDbFile(): string {
  const envPath = process.env.DB_PATH;

  if (envPath) {
    return path.resolve(envPath);
  }

  return resolveDefaultDbFile();
}

const DB_FILE = resolveDbFile();
const DB_BACKUP = DB_FILE + '.bak';
const DATA_DIR = path.dirname(DB_FILE);
// 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 applicationDataSource: DataSource | undefined;
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);
}

function restoreFromBackup(reason: string): Uint8Array | undefined {
  if (!fs.existsSync(DB_BACKUP)) {
    console.error(`[DB] ${reason}. No backup available - starting with a fresh database`);

    return undefined;
  }

  const backup = new Uint8Array(fs.readFileSync(DB_BACKUP));

  if (!isValidSqlite(backup)) {
    console.error(`[DB] ${reason}. Backup is also invalid - starting with a fresh database`);

    return undefined;
  }

  fs.copyFileSync(DB_BACKUP, DB_FILE);
  console.warn('[DB] Restored database from backup', DB_BACKUP);

  return backup;
}

async function migrateLegacyPackagedDatabase(): Promise<void> {
  if (process.env.DB_PATH || !isPackagedRuntime() || path.resolve(DB_FILE) === path.resolve(LEGACY_PACKAGED_DB_FILE)) {
    return;
  }

  let migrated = false;

  if (!fs.existsSync(DB_FILE)) {
    if (fs.existsSync(LEGACY_PACKAGED_DB_FILE)) {
      await fsp.copyFile(LEGACY_PACKAGED_DB_FILE, DB_FILE);
      migrated = true;
    } else if (fs.existsSync(LEGACY_PACKAGED_DB_BACKUP)) {
      await fsp.copyFile(LEGACY_PACKAGED_DB_BACKUP, DB_FILE);
      migrated = true;
    }
  }

  if (!fs.existsSync(DB_BACKUP) && fs.existsSync(LEGACY_PACKAGED_DB_BACKUP)) {
    await fsp.copyFile(LEGACY_PACKAGED_DB_BACKUP, DB_BACKUP);
    migrated = true;
  }

  if (migrated) {
    console.log('[DB] Migrated packaged database files to:', DATA_DIR);
    console.log('[DB] Legacy packaged database location was:', LEGACY_PACKAGED_DB_FILE);
  }
}

/**
 * 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 server loads the database.
 */
function safeguardDbFile(): Uint8Array | undefined {
  if (!fs.existsSync(DB_FILE)) {
    console.warn(`[DB] ${DB_FILE} is missing - checking backup`);

    return restoreFromBackup('Database file missing');
  }

  const data = new Uint8Array(fs.readFileSync(DB_FILE));

  if (isValidSqlite(data)) {
    // Good file - rotate it into the backup slot.
    fs.copyFileSync(DB_FILE, DB_BACKUP);
    console.log('[DB] Backed up database to', DB_BACKUP);

    return data;
  }

  // The main file is corrupt or empty.
  console.warn(`[DB] ${DB_FILE} appears corrupt (${data.length} bytes) - checking backup`);

  return restoreFromBackup(`Database file is invalid (${data.length} bytes)`);
}

function resolveSqlJsConfig(): { locateFile: (file: string) => string } {
  return {
    locateFile: (file) => {
      const bundledBinaryPath = path.join(__dirname, '..', '..', 'node_modules', 'sql.js', 'dist', file);

      return findExistingPath(
        resolveRuntimePath(file),
        bundledBinaryPath
      ) ?? bundledBinaryPath;
    }
  };
}

/**
 * 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, DB_FILE);
      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 = DB_FILE + '.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 function getDataSource(): DataSource {
  if (!applicationDataSource?.isInitialized) {
    throw new Error('DataSource not initialised');
  }

  return applicationDataSource;
}

export async function initDatabase(): Promise<void> {
  if (!fs.existsSync(DATA_DIR))
    fs.mkdirSync(DATA_DIR, { recursive: true });

  await migrateLegacyPackagedDatabase();

  const database = safeguardDbFile();

  try {
    applicationDataSource = new DataSource({
      type: 'sqljs',
      database,
      entities: [
        AuthUserEntity,
        ServerEntity,
        ServerTagEntity,
        ServerChannelEntity,
        ServerRoleEntity,
        ServerUserRoleEntity,
        ServerChannelPermissionEntity,
        JoinRequestEntity,
        ServerMembershipEntity,
        ServerInviteEntity,
        ServerBanEntity,
        GameMatchMissEntity,
        ServerPluginRequirementEntity,
        ServerPluginEventDefinitionEntity,
        PluginDataEntity,
        ServerPluginSettingsEntity,
        PluginUserMetadataEntity
      ],
      migrations: serverMigrations,
      synchronize: process.env.DB_SYNCHRONIZE === 'true',
      logging: false,
      autoSave: true,
      autoSaveCallback: atomicSave,
      sqlJsConfig: resolveSqlJsConfig()
    });
  } catch (error) {
    console.error('[DB] Failed to configure the sql.js data source', error);
    throw error;
  }

  try {
    await applicationDataSource.initialize();
  } catch (error) {
    console.error('[DB] Failed to initialise the sql.js data source', error);
    throw error;
  }

  console.log('[DB] Connection initialised at:', DB_FILE);

  if (process.env.DB_SYNCHRONIZE !== 'true') {
    await applicationDataSource.runMigrations();
    console.log('[DB] Migrations executed');
  } else {
    console.log('[DB] Synchronize mode - migrations skipped');
  }
}

export async function destroyDatabase(): Promise<void> {
  if (applicationDataSource?.isInitialized) {
    await applicationDataSource.destroy();
    applicationDataSource = undefined;
    console.log('[DB] Connection closed');
  }
}