@joystick.js/db-canary 0.0.0-canary.2209

package/src/server/index.js

@@ -0,0 +1,540 @@
+/**
+ * @fileoverview JoystickDB TCP server implementation with clustering, authentication, and replication support.
+ * 
+ * This is the main server entry point that handles TCP connections, message parsing, authentication,
+ * and routing operations to appropriate handlers. Supports clustering, backup/restore, replication,
+ * and comprehensive admin operations.
+ */
+
+import net from 'net';
+import { decode as decode_messagepack } from 'msgpackr';
+import op_types from './lib/op_types.js';
+import safe_json_parse from './lib/safe_json_parse.js';
+import { load_settings, get_settings, get_port_configuration } from './lib/load_settings.js';
+import { send_error } from './lib/send_response.js';
+import { start_cluster } from './cluster/index.js';
+import create_logger from './lib/logger.js';
+import { initialize_database, cleanup_database } from './lib/query_engine.js';
+import { create_message_parser, encode_message } from './lib/tcp_protocol.js';
+import { create_connection_manager } from './lib/connection_manager.js';
+import { shutdown_write_queue } from './lib/write_queue.js';
+import { 
+  setup_authentication, 
+  verify_password, 
+  get_client_ip, 
+  is_rate_limited,
+  initialize_auth_manager,
+  reset_auth_state
+} from './lib/auth_manager.js';
+import {
+  initialize_api_key_manager,
+  verify_user_password,
+  check_admin_user_exists
+} from './lib/api_key_manager.js';
+import { 
+  restore_backup, 
+  start_backup_schedule, 
+  stop_backup_schedule 
+} from './lib/backup_manager.js';
+import { 
+  initialize_replication_manager, 
+  shutdown_replication_manager
+} from './lib/replication_manager.js';
+import { 
+  initialize_write_forwarder, 
+  shutdown_write_forwarder
+} from './lib/write_forwarder.js';
+import { handle_database_operation, handle_admin_operation, handle_ping_operation } from './lib/operation_dispatcher.js';
+import { start_http_server, stop_http_server } from './lib/http_server.js';
+import { 
+  create_recovery_token, 
+  initialize_recovery_manager,
+  reset_recovery_state
+} from './lib/recovery_manager.js';
+
+/** @type {Set<string>} Set of authenticated client socket IDs */
+const authenticated_clients = new Set();
+
+/** @type {Object|null} Connection manager instance */
+let connection_manager = null;
+
+/**
+ * Handles client authentication requests.
+ * @param {net.Socket} socket - Client socket connection
+ * @param {Object} [data={}] - Authentication data containing password
+ * @param {string} data.password - Client password for authentication
+ */
+export const authentication = async (socket, data = {}) => {
+  if (!data?.password) {
+    const response = { ok: 0, error: 'Authentication operation requires password to be set in data.' };
+    const encoded_response = encode_message(response);
+    socket.write(encoded_response);
+    socket.end();
+    return;
+  }
+  
+  try {
+    const client_ip = get_client_ip(socket);
+    
+    if (is_rate_limited(client_ip)) {
+      const response = { ok: 0, error: 'Too many failed attempts. Please try again later.' };
+      const encoded_response = encode_message(response);
+      socket.write(encoded_response);
+      socket.end();
+      return;
+    }
+    
+    const is_valid = await verify_password(data.password, client_ip);
+    
+    if (!is_valid) {
+      const response = { ok: 0, error: 'Authentication failed' };
+      const encoded_response = encode_message(response);
+      socket.write(encoded_response);
+      socket.end();
+      return;
+    }
+    
+    authenticated_clients.add(socket.id);
+    const auth_response = { 
+      ok: 1, 
+      version: "1.0.0",
+      message: 'Authentication successful' 
+    };
+    const encoded_auth_response = encode_message(auth_response);
+    socket.write(encoded_auth_response);
+  } catch (error) {
+    const response = { ok: 0, error: `Authentication error: ${error.message}` };
+    const encoded_response = encode_message(response);
+    socket.write(encoded_response);
+    socket.end();
+  }
+};
+
+/**
+ * Handles initial server setup, generating authentication password.
+ * @param {net.Socket} socket - Client socket connection
+ * @param {Object} [data={}] - Setup data (currently unused)
+ */
+export const setup = async (socket, data = {}) => {
+  try {
+    const password = setup_authentication();
+    
+    // NOTE: Return response format expected by tests.
+    const response = {
+      ok: 1,
+      password,
+      message: 'Authentication setup completed successfully. Save this password - it will not be shown again.'
+    };
+    const encoded_response = encode_message(response);
+    socket.write(encoded_response);
+  } catch (error) {
+    const response = { ok: 0, error: `Setup error: ${error.message}` };
+    const encoded_response = encode_message(response);
+    socket.write(encoded_response);
+  }
+};
+
+/**
+ * Validates if an operation type is supported.
+ * @param {string} [op_type=''] - Operation type to validate
+ * @returns {boolean} True if operation type is valid
+ * @throws {Error} When no operation type is provided
+ */
+export const check_op_type = (op_type = '') => {
+  if (!op_type) {
+    throw new Error('Must pass an op type for operation.');
+  }
+
+  return op_types.includes(op_type);
+};
+
+/**
+ * Parses raw data from various formats (string, Buffer, MessagePack) into JavaScript objects.
+ * @param {string|Buffer|Object} raw_data - Raw data to parse
+ * @returns {Object|null} Parsed data object or null if parsing fails
+ */
+export const parse_data = (raw_data) => {
+  try {
+    if (typeof raw_data === 'string') {
+      const data_as_json = safe_json_parse(raw_data);
+      return data_as_json;
+    } else if (Buffer.isBuffer(raw_data)) {
+      // NOTE: Decode MessagePack buffer and then parse as JSON.
+      const decoded_messagepack = decode_messagepack(raw_data);
+      if (typeof decoded_messagepack === 'string') {
+        return safe_json_parse(decoded_messagepack);
+      }
+      return decoded_messagepack;
+    } else {
+      // NOTE: raw_data is already decoded from MessagePack by the message parser.
+      return raw_data;
+    }
+  } catch (error) {
+    return null;
+  }
+};
+
+/**
+ * Checks if a socket connection is authenticated.
+ * @param {net.Socket} socket - Socket connection to check
+ * @returns {boolean} True if socket is authenticated
+ */
+const check_authentication = (socket) => {
+  return authenticated_clients.has(socket.id);
+};
+
+
+/**
+ * @typedef {Object} ServerOptions
+ * @property {number} [worker_count] - Number of worker processes
+ * @property {number} [port=1983] - Server port
+ * @property {string} [environment] - Environment name for settings file (defaults to NODE_ENV)
+ */
+
+/**
+ * Creates and configures the JoystickDB TCP server.
+ * @returns {Promise<net.Server>} Configured TCP server instance with cleanup method
+ */
+export const create_server = async () => {
+  const { create_context_logger } = create_logger('server');
+  const log = create_context_logger();
+  
+  /*
+    NOTE:
+    Try to load settings, but don't fail if file doesn't exist (for tests).
+  */
+  let settings = null;
+  try {
+    load_settings();
+    settings = get_settings();
+  } catch (error) {
+    /*
+      NOTE:
+      Settings file doesn't exist - this is OK for tests.
+      We'll handle missing settings in individual operations.
+    */
+  }
+  
+  // NOTE: Handle startup restore if configured.
+  if (settings?.restore_from) {
+    try {
+      log.info('Startup restore requested', { backup_filename: settings.restore_from });
+      
+      const restore_result = await restore_backup(settings.restore_from);
+      
+      log.info('Startup restore completed', {
+        backup_filename: settings.restore_from,
+        duration_ms: restore_result.duration_ms
+      });
+      
+      // NOTE: Remove restore_from from settings after successful restore.
+      const updated_settings = { ...settings };
+      delete updated_settings.restore_from;
+      
+      // NOTE: Update the environment variable with the new settings.
+      process.env.JOYSTICK_DB_SETTINGS = JSON.stringify(updated_settings);
+      
+      // NOTE: Reload settings to reflect the change.
+      load_settings();
+      settings = get_settings();
+      
+      log.info('Removed restore_from from settings after successful restore');
+    } catch (restore_error) {
+      log.error('Startup restore failed', {
+        backup_filename: settings.restore_from,
+        error: restore_error.message
+      });
+      
+      // NOTE: Continue with fresh database if restore fails.
+      log.info('Continuing with fresh database after restore failure');
+    }
+  }
+  
+  initialize_database();
+  initialize_auth_manager();
+  initialize_api_key_manager();
+  initialize_recovery_manager();
+  
+  // NOTE: Initialize replication manager for primary nodes.
+  try {
+    initialize_replication_manager();
+    log.info('Replication manager initialized');
+  } catch (replication_error) {
+    log.warn('Failed to initialize replication manager', { error: replication_error.message });
+  }
+  
+  // NOTE: Initialize write forwarder for secondary nodes.
+  try {
+    initialize_write_forwarder();
+    log.info('Write forwarder initialized');
+  } catch (forwarder_error) {
+    log.warn('Failed to initialize write forwarder', { error: forwarder_error.message });
+  }
+  
+  // NOTE: Start backup scheduling if S3 is configured.
+  if (settings?.s3) {
+    try {
+      start_backup_schedule();
+      log.info('Backup scheduling started');
+    } catch (backup_schedule_error) {
+      log.warn('Failed to start backup scheduling', { error: backup_schedule_error.message });
+    }
+  }
+  
+  connection_manager = create_connection_manager({
+    max_connections: 1000,
+    idle_timeout: 10 * 60 * 1000,
+    request_timeout: 5 * 1000
+  });
+  
+  // NOTE: Start HTTP server for setup and recovery operations.
+  let http_server_instance = null;
+  try {
+    const { http_port } = get_port_configuration();
+    http_server_instance = await start_http_server(http_port);
+    if (http_server_instance) {
+      log.info('HTTP server started', { http_port });
+    }
+  } catch (http_error) {
+    log.warn('Failed to start HTTP server', { error: http_error.message });
+  }
+  
+  const server = net.createServer((socket = {}) => {
+    if (!connection_manager.add_connection(socket)) {
+      return;
+    }
+    
+    const message_parser = create_message_parser();
+    
+    socket.on('data', async (raw_data) => {
+      connection_manager.update_activity(socket.id);
+      
+      try {
+        const messages = message_parser.parse_messages(raw_data);
+        
+        for (const message of messages) {
+          // NOTE: Message is already decoded from MessagePack by the message parser.
+          const parsed_data = message;
+          const op_type = parsed_data?.op || null;
+          
+          if (!op_type) {
+            send_error(socket, { message: 'Missing operation type' });
+            continue;
+          }
+          
+          const is_valid_op_type = check_op_type(op_type);
+          
+          if (!is_valid_op_type) {
+            send_error(socket, { message: 'Invalid operation type' });
+            continue;
+          }
+          
+          const request_timeout = connection_manager.create_request_timeout(socket.id, op_type);
+          
+          try {
+            switch(op_type) {
+              case 'authentication':
+                await authentication(socket, parsed_data?.data || {});
+                break;
+                
+              case 'setup':
+                await setup(socket, parsed_data?.data || {});
+                break;
+                
+              case 'insert_one':
+              case 'update_one':
+              case 'delete_one':
+              case 'bulk_write':
+              case 'find_one':
+              case 'find':
+              case 'create_index':
+              case 'drop_index':
+              case 'get_indexes':
+                // NOTE: Use shared operation dispatcher for database operations.
+                await handle_database_operation(socket, op_type, parsed_data?.data || {}, check_authentication, raw_data.length, connection_manager, authenticated_clients);
+                break;
+                
+              case 'ping':
+                handle_ping_operation(socket);
+                break;
+                
+              case 'admin':
+                // NOTE: Use shared operation dispatcher for admin operations.
+                await handle_admin_operation(socket, parsed_data?.data || {}, check_authentication, connection_manager, authenticated_clients);
+                break;
+                
+              case 'reload':
+                if (!check_authentication(socket)) {
+                  send_error(socket, { message: 'Authentication required' });
+                  break;
+                }
+                
+                try {
+                  let old_settings = null;
+                  try {
+                    old_settings = get_settings();
+                  } catch (error) {
+                    // NOTE: Settings not loaded yet, that's OK.
+                  }
+                  
+                  let new_settings = null;
+                  try {
+                    await load_settings();
+                    new_settings = get_settings();
+                  } catch (error) {
+                    // NOTE: Settings file doesn't exist, create default settings for tests.
+                    new_settings = {
+                      port: 1983,
+                      authentication: {}
+                    };
+                  }
+                  
+                  const reload_info = {
+                    ok: 1,
+                    status: 'success',
+                    message: 'Configuration reloaded successfully',
+                    changes: {
+                      port_changed: old_settings ? old_settings.port !== new_settings.port : false,
+                      authentication_changed: old_settings ? old_settings.authentication?.password_hash !== new_settings.authentication?.password_hash : false
+                    },
+                    timestamp: new Date().toISOString()
+                  };
+                  
+                  const encoded_response = encode_message(reload_info);
+                  socket.write(encoded_response);
+                } catch (error) {
+                  const error_response = { ok: 0, error: `Reload operation failed: ${error.message}` };
+                  const encoded_error_response = encode_message(error_response);
+                  socket.write(encoded_error_response);
+                }
+                break;
+                
+              default:
+                send_error(socket, { message: `Operation ${op_type} not implemented` });
+            }
+          } finally {
+            clearTimeout(request_timeout);
+          }
+        }
+      } catch (error) {
+        log.error('Message parsing failed', {
+          client_id: socket.id,
+          error: error.message
+        });
+        
+        send_error(socket, { message: 'Invalid message format' });
+        socket.end();
+      }
+    });
+
+    socket.on('end', () => {
+      log.info('Client disconnected', { socket_id: socket.id });
+      authenticated_clients.delete(socket.id);
+      connection_manager.remove_connection(socket.id);
+    });
+    
+    socket.on('error', (error) => {
+      log.error('Socket error', {
+        socket_id: socket.id,
+        error: error.message
+      });
+      
+      authenticated_clients.delete(socket.id);
+      connection_manager.remove_connection(socket.id);
+    });
+  });
+  
+  // NOTE: Add cleanup method for tests.
+  server.cleanup = async () => {
+    try {
+      // NOTE: Stop HTTP server.
+      await stop_http_server();
+      
+      // NOTE: Stop backup scheduling.
+      stop_backup_schedule();
+      
+      // NOTE: Shutdown replication manager.
+      await shutdown_replication_manager();
+      
+      // NOTE: Shutdown write forwarder.
+      await shutdown_write_forwarder();
+      
+      // NOTE: Shutdown connection manager.
+      if (connection_manager) {
+        connection_manager.shutdown();
+      }
+      
+      // NOTE: Clear authenticated clients.
+      authenticated_clients.clear();
+      
+      // NOTE: Shutdown write queue.
+      await shutdown_write_queue();
+      
+      // NOTE: Wait a bit to ensure all operations are complete.
+      await new Promise(resolve => setTimeout(resolve, 100));
+      
+      // NOTE: Cleanup database (after everything else is shut down).
+      await cleanup_database();
+      
+      // NOTE: Reset auth state.
+      reset_auth_state();
+      
+      // NOTE: Reset recovery state.
+      reset_recovery_state();
+    } catch (error) {
+      // NOTE: Ignore cleanup errors in tests.
+    }
+  };
+  
+  return server;
+};
+
+if (import.meta.url === `file://${process.argv[1]}`) {
+  const { create_context_logger } = create_logger('main');
+  const log = create_context_logger();
+  
+  // NOTE: Check for recovery token generation flag.
+  if (process.argv.includes('--generate-recovery-token')) {
+    try {
+      initialize_recovery_manager();
+      const recovery_info = create_recovery_token();
+      
+      console.log('Emergency Recovery Token Generated');
+      console.log(`Visit: ${recovery_info.url}`);
+      console.log('Token expires in 10 minutes');
+      
+      log.info('Recovery token generated via CLI', {
+        expires_at: new Date(recovery_info.expires_at).toISOString()
+      });
+      
+      process.exit(0);
+    } catch (error) {
+      console.error('Failed to generate recovery token:', error.message);
+      log.error('Recovery token generation failed', { error: error.message });
+      process.exit(1);
+    }
+  }
+  
+  // NOTE: Get port configuration from settings or use defaults.
+  const { tcp_port, http_port } = get_port_configuration();
+  
+  const options = {
+    worker_count: process.env.WORKER_COUNT ? parseInt(process.env.WORKER_COUNT) : undefined,
+    port: tcp_port,
+    environment: process.env.NODE_ENV || 'development'
+  };
+  
+  const { has_settings } = await import('./lib/load_settings.js');
+  const has_db_settings = has_settings();
+  
+  log.info('Starting JoystickDB server...', {
+    workers: options.worker_count || 'auto',
+    tcp_port: tcp_port,
+    http_port: http_port,
+    environment: options.environment,
+    has_settings: has_db_settings,
+    port_source: has_db_settings ? 'JOYSTICK_DB_SETTINGS' : 'default'
+  });
+  
+  start_cluster(options);
+}