Checkpoint/index.js

import { mkdir, readFile } from 'fs/promises';
import { writeFileSync, readFileSync, unlinkSync, existsSync } from 'fs';
import { join, dirname } from 'path';
import { fileURLToPath } from 'url';
import { secureImportModule } from './utils/plugins.js';
import * as logs from './utils/logs.js';
import express from 'express';
import { createServer } from 'http';
import { WebSocketServer } from 'ws';
import { spawn } from 'child_process';

// Load environment variables from .env file
import dotenv from 'dotenv';
dotenv.config();

// Stop daemon: if run with -k, kill the running process and exit.
if (process.argv.includes('-k')) {
  const pidFile = join(dirname(fileURLToPath(import.meta.url)), 'checkpoint.pid');
  if (existsSync(pidFile)) {
    const pid = parseInt(readFileSync(pidFile, 'utf8'), 10);
    try {
      process.kill(pid);
      unlinkSync(pidFile);
      console.log(`Stopped daemon (pid ${pid})`);
    } catch (err) {
      console.error(`Failed to stop pid ${pid}: ${err}`);
    }
  } else {
    console.error(`No pid file found at ${pidFile}`);
  }
  process.exit(0);
}

// Daemonize: if run with -d, kill any existing daemon, then re-spawn detached, write pid file, and exit parent.
if (process.argv.includes('-d')) {
  const pidFile = join(dirname(fileURLToPath(import.meta.url)), 'checkpoint.pid');
  // If already running, stop the old daemon
  if (existsSync(pidFile)) {
    const oldPid = parseInt(readFileSync(pidFile, 'utf8'), 10);
    try {
      process.kill(oldPid);
      console.log(`Stopped old daemon (pid ${oldPid})`);
    } catch (e) {
      console.error(`Failed to stop old daemon (pid ${oldPid}): ${e}`);
    }
    try {
      unlinkSync(pidFile);
    } catch {}
  }
  // Spawn new background process
  const args = process.argv.slice(1).filter((arg) => arg !== '-d');
  const cp = spawn(process.argv[0], args, {
    detached: true,
    stdio: 'ignore'
  });
  cp.unref();
  writeFileSync(pidFile, cp.pid.toString(), 'utf8');
  console.log(`Daemonized (pid ${cp.pid}), pid stored in ${pidFile}`);
  process.exit(0);
}

const pluginRegistry = [];
export function registerPlugin(pluginName, handler) {
  pluginRegistry.push({ name: pluginName, handler });
}
/**
 * Return the array of middleware handlers in registration order.
 */
export function loadPlugins() {
  return pluginRegistry.map((item) => item.handler);
}
/**
 * Return the names of all registered plugins.
 */
export function getRegisteredPluginNames() {
  return pluginRegistry.map((item) => item.name);
}
/**
 * Freeze plugin registry to prevent further registration and log the final set.
 */
export function freezePlugins() {
  Object.freeze(pluginRegistry);
  pluginRegistry.forEach((item) => Object.freeze(item));
  logs.msg('Plugin registration frozen');
}

// Determine root directory for config loading
const __dirname = dirname(fileURLToPath(import.meta.url));
export const rootDir = __dirname;

export async function loadConfig(name, target) {
  const configPath = join(rootDir, 'config', `${name}.toml`);
  const txt = await readFile(configPath, 'utf8');
  const { default: toml } = await import('@iarna/toml');
  Object.assign(target, toml.parse(txt));
  logs.config(name, 'loaded');
}

async function initDataDirectories() {
  logs.section('INIT');
  const directories = [join(rootDir, 'data'), join(rootDir, 'db'), join(rootDir, 'config')];
  for (const dirPath of directories) {
    try {
      await mkdir(dirPath, { recursive: true });
    } catch {}
  }
  logs.init('Data directories are now in place');
}

function staticFileMiddleware() {
  const router = express.Router();
  router.use('/webfont', express.static(join(rootDir, 'pages/interstitial/webfont'), {
    maxAge: '7d'
  }));
  router.use('/js', express.static(join(rootDir, 'pages/interstitial/js'), {
    maxAge: '7d'
  }));
  return router;
}

async function main() {
  await initDataDirectories();

  logs.section('CONFIG');
  logs.config('checkpoint', 'loaded');
  logs.config('ipfilter', 'loaded');
  logs.config('proxy', 'loaded');
  logs.config('stats', 'loaded');

  logs.section('OPERATIONS');
  
  const app = express();
  const server = createServer(app);
  
  // Trust proxy headers (important for proper protocol detection)
  app.set('trust proxy', true);
  
  // Initialize WebSocket server
  const wss = new WebSocketServer({ noServer: true });
  
  // Store WebSocket handlers
  let wsHandlers = {};
  
  try {
    await secureImportModule('checkpoint.js');
  } catch (e) {
    logs.error('checkpoint', `Failed to load checkpoint plugin: ${e}`);
  }
  try {
    await secureImportModule('plugins/ipfilter.js');
  } catch (e) {
    logs.error('ipfilter', `Failed to load IP filter plugin: ${e}`);
  }
  try {
    await secureImportModule('plugins/proxy.js');
    const mod = await import('./plugins/proxy.js');
    if (mod.proxyWebSocketHandler) {
      wsHandlers = mod.proxyWebSocketHandler;
    }
  } catch (e) {
    logs.error('proxy', `Failed to load proxy plugin: ${e}`);
  }
  try {
    await secureImportModule('plugins/stats.js');
  } catch (e) {
    logs.error('stats', `Failed to load stats plugin: ${e}`);
  }

  // Register static middleware
  app.use(staticFileMiddleware());

  logs.section('PLUGINS');
  // Ensure ipfilter runs first by moving it to front of the registry
  const ipIndex = pluginRegistry.findIndex((item) => item.name === 'ipfilter');
  if (ipIndex > 0) {
    const [ipEntry] = pluginRegistry.splice(ipIndex, 1);
    pluginRegistry.unshift(ipEntry);
  }
  pluginRegistry.forEach((item) => logs.msg(item.name));
  logs.section('SYSTEM');
  freezePlugins();

  // Apply all plugin middlewares to Express
  const middlewareHandlers = loadPlugins();
  middlewareHandlers.forEach(handler => {
    if (typeof handler === 'function') {
      // Wrap plugin handlers to work with Express
      app.use(async (req, res, next) => {
        try {
          const result = await handler(req, { upgrade: () => false });
          if (result instanceof Response) {
            // Convert Response to Express response
            res.status(result.status);
            result.headers.forEach((value, key) => {
              res.setHeader(key, value);
            });
            const body = await result.text();
            res.send(body);
          } else {
            next();
          }
        } catch (err) {
          logs.error('server', `Handler error: ${err}`);
          next(err);
        }
      });
    } else if (handler && handler.middleware) {
      // If plugin exports Express middleware directly
      app.use(handler.middleware);
    }
  });

  // 404 handler
  app.use((req, res) => {
    res.status(404).send('Not Found');
  });

  // Error handler
  app.use((err, req, res, next) => {
    logs.error('server', `Server error: ${err.message}`);
    res.status(500).send(`Server Error: ${err.message}`);
  });

  // Handle WebSocket upgrades
  server.on('upgrade', (request, socket, head) => {
    wss.handleUpgrade(request, socket, head, (ws) => {
      wss.emit('connection', ws, request);
    });
  });

  // WebSocket connection handler
  if (wsHandlers.open) {
    wss.on('connection', (ws, request) => {
      ws.data = {};
      if (wsHandlers.open) wsHandlers.open(ws);
      if (wsHandlers.message) {
        ws.on('message', (message) => wsHandlers.message(ws, message));
      }
      if (wsHandlers.close) {
        ws.on('close', (code, reason) => wsHandlers.close(ws, code, reason));
      }
      if (wsHandlers.error) {
        ws.on('error', (err) => wsHandlers.error(ws, err));
      }
    });
  }

  logs.section('SERVER');
  const portNumber = Number(process.env.PORT || 3000);
  server.listen(portNumber, () => {
    logs.server(`🚀 Server is up and running on port ${portNumber}...`);
    logs.section('REQ LOGS');
  });
}

main();