main.cpp

Go to the documentation of this file.

/// \file main.cpp
/// \brief Точка входа: event loop, обработка сигналов, метрики Prometheus.
/// \details Содержит главный цикл приложения: загрузка конфигурации,
///          инициализация БД и XDP, периодический опрос списков и
///          обслуживание HTTP-эндпоинта /metrics.
 
#include "main.h"
#include "db.h"
#include "xdp_loader.h"
 
#include <arpa/inet.h>
#include <csignal>
#include <cstring>
#include <ctime>
#include <fstream>
#include <getopt.h>
#include <iostream>
#include <map>
#include <shared_mutex>
#include <sstream>
#include <thread>
#include <unistd.h>
 
#include "httplib.h"
 
/**********************************************************************/
/*  Global state                                                      */
/**********************************************************************/
 
std::atomic<bool> g_running{true};
std::atomic<bool> g_xdp_attached{false};
std::atomic<bool> g_detach_requested{false};
 
/// Потокобезопасный аккумулятор метрик Prometheus.
struct MetricsStore {
  std::map<uint32_t, uint64_t>
      total_drops; ///< Суммарное количество дропов с запуска процесса.
  std::map<uint32_t, uint64_t>
      interval_drops; ///< Дропов за последний интервал опроса.
  std::map<uint32_t, uint64_t> hourly_drops; ///< Дропов с начала текущего часа.
  std::map<uint32_t, uint64_t> daily_drops; ///< Дропов с начала текущего дня.
  size_t whitelist_size = 0; ///< Текущее количество записей в whitelist.
  size_t blacklist_size = 0; ///< Текущее количество записей в blacklist.
  bool xdp_attached = false; ///< Прикреплён ли XDP в данный момент.
  std::shared_mutex mtx; ///< Защищает все поля выше.
};
 
static XdpBlocker *g_blocker = nullptr;
 
/**********************************************************************/
/*  Signal handler                                                    */
/**********************************************************************/
 
/// Обрабатывает сигналы ОС для управления жизненным циклом.
///
/// - SIGINT / SIGTERM: запрос плавного завершения.
/// - SIGUSR1:          запрос открепления XDP (аварийное отключение).
/// @param[in] sig  Номер сигнала.
static void HandleSignal(int sig) {
  switch (sig) {
  case SIGINT:
  case SIGTERM:
    g_running = false;
    break;
  case SIGUSR1:
    g_detach_requested = true;
    break;
  }
}
 
/**********************************************************************/
/*  Argument parsing                                                  */
/**********************************************************************/
 
/// Парсит аргументы командной строки через getopt_long.
/// @param[in] argc  Количество аргументов из main().
/// @param[in] argv  Массив аргументов из main().
/// @return          Заполненная структура Config (для пропущенных флагов —
/// значения по умолчанию).
static Config ParseArgs(int argc, char *argv[]) {
  Config cfg;
  static struct option long_opts[] = {
      {"iface", required_argument, nullptr, 'i'},
      {"db-path", required_argument, nullptr, 'd'},
      {"whitelist", required_argument, nullptr, 'w'},
      {"interval", required_argument, nullptr, 't'},
      {"metrics-listen", required_argument, nullptr, 'm'},
      {"help", no_argument, nullptr, 'h'},
      {nullptr, 0, nullptr, 0}};
 
  int opt;
  while ((opt = getopt_long(argc, argv, "i:d:w:t:m:h", long_opts, nullptr)) !=
         -1) {
    switch (opt) {
    case 'i':
      cfg.iface = optarg;
      break;
    case 'd':
      cfg.db_path = optarg;
      break;
    case 'w':
      cfg.whitelist_conf = optarg;
      break;
    case 't':
      cfg.interval_sec = std::stoi(optarg);
      break;
    case 'm':
      cfg.metrics_listen = optarg;
      break;
    case 'h':
      std::cout
          << "Usage: ebpf-blocker [options]\n"
          << "  -i, --iface IFACE          Network interface (default: eth0)\n"
          << "  -d, --db-path PATH          SQLite database path\n"
          << "  -w, --whitelist PATH        Whitelist config file\n"
          << "  -t, --interval SEC          Poll interval (default: 60)\n"
          << "  -m, --metrics-listen ADDR   Metrics HTTP listen address "
             "(default: 0.0.0.0:9101)\n"
          << "  -h, --help                  Show this help\n";
      exit(0);
    }
  }
  return cfg;
}
 
/**********************************************************************/
/*  Whitelist file parser                                             */
/**********************************************************************/
 
/// Читает записи белого списка из конфигурационного файла.
///
/// Формат: одна IP/CIDR на строку, опционально '# комментарий' в конце.
/// Пустые строки и строки, начинающиеся с '#', игнорируются.
/// @param[in] path  Путь к файлу конфигурации белого списка.
/// @return          Список распарсенных записей.
static std::vector<BlockEntry> ReadWhitelistFile(const std::string &path) {
  std::vector<BlockEntry> entries;
  std::ifstream file(path);
  if (!file.is_open())
    return entries;
 
  std::string line;
  while (std::getline(file, line)) {
    if (line.empty() || line[0] == '#')
      continue;
    auto comment_pos = line.find('#');
    std::string ip =
        (comment_pos != std::string::npos) ? line.substr(0, comment_pos) : line;
    ip.erase(ip.find_last_not_of(" \t\r\n") + 1);
    ip.erase(0, ip.find_first_not_of(" \t\r\n"));
    if (ip.empty())
      continue;
    std::string comment = "whitelist.conf";
    if (comment_pos != std::string::npos) {
      comment = line.substr(comment_pos + 1);
      comment.erase(0, comment.find_first_not_of(" \t\r\n"));
    }
    entries.push_back({ip, comment});
  }
  return entries;
}
 
#include <algorithm>
 
/**********************************************************************/
/*  Logging helpers                                                    */
/**********************************************************************/
 
/// Логирует текущее количество записей в whitelist и blacklist (heartbeat).
/// @param[in] whitelist  Список записей белого списка.
/// @param[in] blacklist  Список записей чёрного списка.
static void PrintStats(const std::vector<BlockEntry> &whitelist,
                       const std::vector<BlockEntry> &blacklist) {
  std::cout << "whitelist: " << whitelist.size() << " entries, "
            << "blacklist: " << blacklist.size() << " entries" << std::endl;
}
 
/// Удобный алиас для карты соответствия IP (uint32) количеству дропов.
using DropMap = std::map<uint32_t, uint64_t>;
 
/// Логирует статистику дропов за последний интервал опроса (всего + топ).
/// @param[in] drops  Вектор пар (source_ip, количество_дропов).
static void
PrintIntervalDrops(const std::vector<std::pair<uint32_t, uint64_t>> &drops) {
  uint64_t total = 0;
  for (const auto &d : drops)
    total += d.second;
 
  std::cout << "drops: " << total << " packets";
 
  if (!drops.empty()) {
    auto top = std::max_element(
        drops.begin(), drops.end(),
        [](const auto &a, const auto &b) { return a.second < b.second; });
 
    char buf[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &top->first, buf, sizeof(buf));
    std::cout << " | top: " << buf << " (" << top->second << ")";
  }
 
  std::cout << std::endl;
}
 
/// Логирует периодическую статистику дропов (почасовую/ежедневную) с топ-5
/// нарушителями.
/// @param[in] label      Метка периода ("почасовая" или "ежедневная").
/// @param[in] drops      Карта дропов (IP -> количество).
/// @param[in] timestamp  Временная метка для лога.
static void PrintTopDrops(const std::string &label, const DropMap &drops,
                          const char *timestamp) {
  uint64_t total = 0;
  for (const auto &d : drops)
    total += d.second;
 
  std::cout << "[" << timestamp << "] === " << label << " ===" << std::endl;
  std::cout << "  drops: " << total << " packets" << std::endl;
 
  if (!drops.empty()) {
    using Pair = std::pair<uint32_t, uint64_t>;
    std::vector<Pair> sorted(drops.begin(), drops.end());
    std::sort(sorted.begin(), sorted.end(),
              [](const Pair &a, const Pair &b) { return a.second > b.second; });
 
    std::cout << "  top offenders:" << std::endl;
    int count = 0;
    char buf[INET_ADDRSTRLEN];
    for (const auto &d : sorted) {
      if (count++ >= 5)
        break;
      inet_ntop(AF_INET, &d.first, buf, sizeof(buf));
      std::cout << "    " << buf << "  " << d.second << std::endl;
    }
  }
  std::cout << std::endl;
}
 
/**********************************************************************/
/*  Main entry point                                                  */
/**********************************************************************/
 
/// Точка входа.
///
/// Инициализирует БД, загружает XDP-программу, запускает event loop
/// с периодическим опросом списков и HTTP-сервером метрик Prometheus.
/// @return 0 при успешном завершении.
int main(int argc, char *argv[]) {
  Config cfg = ParseArgs(argc, argv);
 
  signal(SIGINT, HandleSignal);
  signal(SIGTERM, HandleSignal);
  signal(SIGUSR1, HandleSignal);
 
  try {
    BlockerDB db(cfg.db_path);
    db.Migrate();
 
    XdpBlocker blocker;
    g_blocker = &blocker;
 
    blocker.Load(cfg.iface);
    g_xdp_attached = true;
 
    auto whitelist_conf = ReadWhitelistFile(cfg.whitelist_conf);
    if (!whitelist_conf.empty()) {
      std::cout << "loaded " << whitelist_conf.size() << " entries from "
                << cfg.whitelist_conf << std::endl;
    }
 
    auto whitelist_db = db.GetWhitelist();
    whitelist_conf.insert(whitelist_conf.end(), whitelist_db.begin(),
                          whitelist_db.end());
 
    auto blacklist = db.GetBlacklist();
 
    if (whitelist_conf.empty()) {
      std::cerr << "WARNING: whitelist is empty! "
                << "Your SSH connection may be at risk. "
                << "Sleeping 30 seconds before applying rules..." << std::endl;
      for (int i = 0; i < 30 && g_running; i++)
        sleep(1);
    }
 
    blocker.UpdateWhitelist(whitelist_conf);
    blocker.UpdateBlacklist(blacklist);
 
    std::cout << "blocker started on " << cfg.iface << " (poll every "
              << cfg.interval_sec << "s)" << std::endl;
    PrintStats(whitelist_conf, blacklist);
 
    auto metrics = std::make_shared<MetricsStore>();
    {
      std::lock_guard<std::shared_mutex> lock(metrics->mtx);
      metrics->whitelist_size = whitelist_conf.size();
      metrics->blacklist_size = blacklist.size();
      metrics->xdp_attached = true;
      metrics->total_drops.clear();
    }
 
    std::thread http_thread([metrics, &cfg]() {
      httplib::Server svr;
      svr.Get("/metrics", [metrics](const httplib::Request &,
                                    httplib::Response &res) {
        std::shared_lock<std::shared_mutex> lock(metrics->mtx);
        std::ostringstream ss;
        ss << "# HELP ebpf_blocker_whitelist_entries Current whitelist "
              "entries\n"
           << "# TYPE ebpf_blocker_whitelist_entries gauge\n"
           << "ebpf_blocker_whitelist_entries " << metrics->whitelist_size
           << "\n"
           << "# HELP ebpf_blocker_blacklist_entries Current blacklist "
              "entries\n"
           << "# TYPE ebpf_blocker_blacklist_entries gauge\n"
           << "ebpf_blocker_blacklist_entries " << metrics->blacklist_size
           << "\n"
           << "# HELP ebpf_blocker_xdp_attached XDP attached flag (1=attached, "
              "0=detached)\n"
           << "# TYPE ebpf_blocker_xdp_attached gauge\n"
           << "ebpf_blocker_xdp_attached " << (metrics->xdp_attached ? 1 : 0)
           << "\n"
           << "# HELP ebpf_blocker_drops_total Total dropped packets since "
              "process start\n"
           << "# TYPE ebpf_blocker_drops_total counter\n";
        char ip_buf[INET_ADDRSTRLEN];
        for (const auto &[ip, count] : metrics->total_drops) {
          inet_ntop(AF_INET, &ip, ip_buf, sizeof(ip_buf));
          ss << "ebpf_blocker_drops_total{ip=\"" << ip_buf << "\"} " << count
             << "\n";
        }
        ss << "# HELP ebpf_blocker_drops_interval Drops in last poll interval\n"
           << "# TYPE ebpf_blocker_drops_interval gauge\n";
        if (metrics->interval_drops.empty()) {
          ss << "ebpf_blocker_drops_interval 0\n";
        } else {
          for (const auto &[ip, count] : metrics->interval_drops) {
            inet_ntop(AF_INET, &ip, ip_buf, sizeof(ip_buf));
            ss << "ebpf_blocker_drops_interval{ip=\"" << ip_buf << "\"} "
               << count << "\n";
          }
        }
        ss << "# HELP ebpf_blocker_drops_hourly Drops since start of current "
              "hour\n"
           << "# TYPE ebpf_blocker_drops_hourly gauge\n";
        if (metrics->hourly_drops.empty()) {
          ss << "ebpf_blocker_drops_hourly 0\n";
        } else {
          for (const auto &[ip, count] : metrics->hourly_drops) {
            inet_ntop(AF_INET, &ip, ip_buf, sizeof(ip_buf));
            ss << "ebpf_blocker_drops_hourly{ip=\"" << ip_buf << "\"} " << count
               << "\n";
          }
        }
        ss << "# HELP ebpf_blocker_drops_daily Drops since start of current "
              "day\n"
           << "# TYPE ebpf_blocker_drops_daily gauge\n";
        if (metrics->daily_drops.empty()) {
          ss << "ebpf_blocker_drops_daily 0\n";
        } else {
          for (const auto &[ip, count] : metrics->daily_drops) {
            inet_ntop(AF_INET, &ip, ip_buf, sizeof(ip_buf));
            ss << "ebpf_blocker_drops_daily{ip=\"" << ip_buf << "\"} " << count
               << "\n";
          }
        }
        res.set_content(ss.str(), "text/plain; charset=utf-8");
      });
      std::cout << "metrics HTTP server listening on " << cfg.metrics_listen
                << std::endl;
      auto colon = cfg.metrics_listen.find_last_of(':');
      auto host = cfg.metrics_listen.substr(0, colon);
      auto port = std::stoi(cfg.metrics_listen.substr(colon + 1));
      if (!svr.listen(host, port))
        std::cerr << "failed to start metrics HTTP server on "
                  << cfg.metrics_listen << std::endl;
    });
    http_thread.detach();
 
    int prev_hour = -1;
    int prev_day = -1;
 
    while (g_running) {
      for (int i = 0; i < cfg.interval_sec && g_running; i++)
        sleep(1);
 
      if (!g_running)
        break;
 
      if (g_detach_requested) {
        g_blocker->Unload();
        g_xdp_attached = false;
        g_detach_requested = false;
        std::cout << "SIGUSR1: XDP detached, traffic passes through"
                  << std::endl;
      }
 
      time_t now = time(nullptr);
      struct tm tm_now;
      localtime_r(&now, &tm_now);
 
      int cur_hour = tm_now.tm_hour;
      int cur_min = tm_now.tm_min;
      int cur_day = tm_now.tm_yday;
 
      if (!g_xdp_attached) {
        try {
          blocker.Load(cfg.iface);
          g_xdp_attached = true;
          std::cout << "XDP re-attached" << std::endl;
        } catch (const std::exception &e) {
          std::cerr << "failed to re-attach XDP: " << e.what() << std::endl;
          continue;
        }
      }
 
      try {
        whitelist_conf = ReadWhitelistFile(cfg.whitelist_conf);
        whitelist_db = db.GetWhitelist();
        whitelist_conf.insert(whitelist_conf.end(), whitelist_db.begin(),
                              whitelist_db.end());
 
        blacklist = db.GetBlacklist();
 
        blocker.UpdateWhitelist(whitelist_conf);
        blocker.UpdateBlacklist(blacklist);
 
        auto drops = blocker.ReadAndClearDropCounters();
 
        for (const auto &d : drops)
          db.UpdateLastBlockedAt(d.first);
 
        PrintStats(whitelist_conf, blacklist);
        PrintIntervalDrops(drops);
 
        {
          std::lock_guard<std::shared_mutex> lock(metrics->mtx);
          for (const auto &d : drops) {
            metrics->total_drops[d.first] += d.second;
            metrics->hourly_drops[d.first] += d.second;
            metrics->daily_drops[d.first] += d.second;
          }
          metrics->interval_drops.clear();
          for (const auto &d : drops)
            metrics->interval_drops[d.first] += d.second;
          metrics->whitelist_size = whitelist_conf.size();
          metrics->blacklist_size = blacklist.size();
          metrics->xdp_attached = g_xdp_attached;
        }
 
        char ts[64];
        strftime(ts, sizeof(ts), "%Y-%m-%d %H:%M", &tm_now);
 
        if (cur_min == 0 && cur_hour != prev_hour) {
          DropMap snapshot;
          {
            std::lock_guard<std::shared_mutex> lock(metrics->mtx);
            snapshot = metrics->hourly_drops;
            metrics->hourly_drops.clear();
          }
          if (!snapshot.empty())
            PrintTopDrops("Hourly stats (last 60 min)", snapshot, ts);
          prev_hour = cur_hour;
        }
 
        if (cur_hour == 0 && cur_min == 0 && cur_day != prev_day) {
          DropMap snapshot;
          {
            std::lock_guard<std::shared_mutex> lock(metrics->mtx);
            snapshot = metrics->daily_drops;
            metrics->daily_drops.clear();
          }
          if (!snapshot.empty())
            PrintTopDrops("Daily stats (last 24h)", snapshot, ts);
          prev_day = cur_day;
        }
      } catch (const std::exception &e) {
        std::cerr << "update error: " << e.what() << std::endl;
      }
    }
 
    std::cout << "shutting down..." << std::endl;
    blocker.Unload();
    g_blocker = nullptr;
 
  } catch (const std::exception &e) {
    std::cerr << "fatal error: " << e.what() << std::endl;
    return 1;
  }
 
  return 0;
}