src/main.cpp - master - smychkov/SStorage

src/main.cpp
220 строк9 KB
Андрей Смычков
security: комплексные фиксы по результатам аудита
25 апр 2026, 10:00
25 апр 2026, 10:005f1be8e
//============================================================================
// SStorage — main entry point
//============================================================================
// Режимы запуска:
//   interactive — REPL (put/get/del/scan/...)
//   http         — HTTP-сервер
//   grpc         — gRPC-сервер
//   both        — HTTP + gRPC
//============================================================================

#include "core/database.hpp"
#include "core/config.hpp"
#include "cli/command.hpp"
#include "server/server.hpp"
#include "util/utils.hpp"

#include <atomic>
#include <chrono>
#include <csignal>
#include <cstdlib>
#include <getopt.h>
#include <iostream>
#include <memory>
#include <string>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>

namespace sstorage {

//============================================================================
// Global pointers для graceful shutdown через signal handler
//============================================================================
static Database* g_db = nullptr;
static Server* g_server = nullptr;
static std::atomic<int> g_shuttingDown{0};

static void safeShutdown() {
    if (g_shuttingDown.exchange(1) != 0) return;
    if (g_server) {
        g_server->stop();
        delete g_server;
        g_server = nullptr;
    }
    if (g_db) {
        std::cerr << "Saving data before exit...\n";
        g_db->close();
        delete g_db;
        g_db = nullptr;
    }
}

static void signalHandler(int signum) {
    safeShutdown();
    std::_Exit(signum);
}

//============================================================================
// Определение директории данных рядом с исполняемым файлом
//============================================================================
static std::string resolveDataDir(const char* argv0) {
    if (!argv0 || !*argv0) return "./data";
    std::string path(argv0);
    auto slash = path.rfind('/');
    if (slash == std::string::npos) return "./data";
    std::string appDir = path.substr(0, slash);
    std::string dataDir = appDir + "/data";
    return dataDir;
}

//============================================================================
// Help
//============================================================================
static void printUsage(const char* progName) {
    std::cout <<
        "Usage: " << progName << " [options]\n"
        "\n"
        "Options:\n"
        "  --data-dir DIR           Data directory (default: ./data near binary)\n"
        "  --memtable-size-mb N     MemTable size (default: 4)\n"
        "  --block-size-kb N        SSTable block size (default: 4)\n"
        "  --block-cache-mb N       Block cache size (default: 64)\n"
        "  --l1-size-mb N           L1 level size target (default: 10)\n"
        "  --level-ratio N          Size ratio between levels (default: 10)\n"
        "  --bloom-bits-per-key N   Bloom filter bits per key (default: 10)\n"
        "  --max-key-size N         Max key size in bytes (default: 4096)\n"
        "  --max-value-size N       Max value size in bytes (default: 1048576)\n"
        "  --http-port N            HTTP port (default: 8080)\n"
        "  --grpc-port N            gRPC port (default: 8081)\n"
        "  --mode MODE              interactive | http | grpc | both (default: interactive)\n"
        "  -h, --help               Show this help\n"
        "\n"
        "Examples:\n"
        "  " << progName << "                      # Interactive CLI\n"
        "  " << progName << " --mode http          # HTTP server only\n"
        "  " << progName << " --mode both          # HTTP + gRPC\n";
}

}

//============================================================================
// main
//============================================================================
int main(int argc, char* argv[]) {
    std::signal(SIGINT, sstorage::signalHandler);
    std::signal(SIGTERM, sstorage::signalHandler);

    sstorage::Config cfg;
    cfg.setDataDirectory(sstorage::resolveDataDir(argv[0]));

    //------------------------------------------------------------------------
    // Парсинг аргументов
    //------------------------------------------------------------------------
    static struct option longOptions[] = {
        {"data-dir", required_argument, 0, 'd'},
        {"memtable-size-mb", required_argument, 0, 'm'},
        {"block-size-kb", required_argument, 0, 'b'},
        {"block-cache-mb", required_argument, 0, 'c'},
        {"l1-size-mb", required_argument, 0, 'l'},
        {"level-ratio", required_argument, 0, 'r'},
        {"bloom-bits-per-key", required_argument, 0, 'B'},
        {"max-key-size", required_argument, 0, 'k'},
        {"max-value-size", required_argument, 0, 'v'},
        {"http-port", required_argument, 0, 'H'},
        {"grpc-port", required_argument, 0, 'G'},
        {"mode", required_argument, 0, 'M'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int opt;
    int idx = 0;
    while ((opt = getopt_long(argc, argv, "d:m:b:c:l:r:B:k:v:H:G:M:h",
                              longOptions, &idx)) != -1) {
        try {
            switch (opt) {
                case 'd': cfg.setDataDirectory(optarg); break;
                case 'm': cfg.lsmOptions().memtableSizeBytes =
                              std::stoul(optarg) * 1024 * 1024; break;
                case 'b': cfg.lsmOptions().blockSize =
                              std::stoul(optarg) * 1024; break;
                case 'c': cfg.setBlockCacheBytes(std::stoul(optarg) * 1024 * 1024); break;
                case 'l': cfg.lsmOptions().l1SizeBytes =
                              std::stoul(optarg) * 1024 * 1024; break;
                case 'r': cfg.lsmOptions().levelRatio = std::stoul(optarg); break;
                case 'B': cfg.lsmOptions().bloomBitsPerKey = std::stoul(optarg); break;
                case 'k': cfg.setMaxKeySize(std::stoul(optarg)); break;
                case 'v': cfg.setMaxValueSize(std::stoul(optarg)); break;
                case 'H': cfg.setHttpPort(
                              static_cast<uint16_t>(std::stoul(optarg))); break;
                case 'G': cfg.setGrpcPort(
                              static_cast<uint16_t>(std::stoul(optarg))); break;
                case 'M': cfg.setMode(optarg); break;
                case 'h': sstorage::printUsage(argv[0]); return 0;
                default: sstorage::printUsage(argv[0]); return 1;
            }
        } catch (const std::exception&) {
            std::cerr << "Error: invalid value for option\n";
            return 1;
        }
    }

    // Валидация mode
    const std::string mode = cfg.mode();
    if (mode != "interactive" && mode != "http" && mode != "grpc" && mode != "both") {
        std::cerr << "Error: invalid mode: " << mode << "\n";
        return 1;
    }

    //------------------------------------------------------------------------
    // Валидация директории данных (защита от path traversal и системных путей)
    //------------------------------------------------------------------------
    {
        std::string resolved;
        if (!sstorage::util::validateDataDirectory(cfg.dataDirectory(), resolved)) {
            std::cerr << "Error: invalid or forbidden data directory: "
                      << cfg.dataDirectory() << "\n";
            return 1;
        }
        cfg.setDataDirectory(resolved);
    }

    //------------------------------------------------------------------------
    // Создание и открытие Database
    //------------------------------------------------------------------------
    sstorage::g_db = new sstorage::Database(cfg);
    if (!sstorage::g_db->open()) {
        std::cerr << "Failed to open database at " << cfg.dataDirectory() << "\n";
        delete sstorage::g_db;
        return 1;
    }

    //------------------------------------------------------------------------
    // Выбор режима работы
    //------------------------------------------------------------------------
    if (mode == "interactive") {
        int result = sstorage::runInteractive(*sstorage::g_db);
        sstorage::g_db->close();
        delete sstorage::g_db;
        sstorage::g_db = nullptr;
        return result;
    }

    // Сетевой режим (http/grpc/both)
    sstorage::g_server = new sstorage::Server(*sstorage::g_db,
                                              cfg.httpPort(), cfg.grpcPort(), mode);
    if (!sstorage::g_server->start()) {
        std::cerr << "Failed to start server in mode: " << mode << "\n";
        delete sstorage::g_server;
        delete sstorage::g_db;
        return 1;
    }

    std::cout << "Server running in " << mode << " mode. Press Ctrl+C to stop.\n";
    while (!sstorage::g_shuttingDown.load()) {
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    sstorage::safeShutdown();
    return 0;
}