// Egress tests: the direct backend end to end, and the switch controller
// against a fake backend.
//
// The switch controller is the subtlest code in the program and the hardest to
// observe in production -- it fires rarely, on a schedule set by a volunteer
// server going bad. Testing it against a real VPNGate node would measure the
// weather. So EgressManager takes an egress factory (egress_manager.h), and
// everything below drives the state machine through a fake whose failures are
// exactly reproducible.
//
// The *selector* underneath is real, not faked: node lists, scoring, and the
// latency prober all run for real, aimed at a loopback listener this file
// starts. That keeps the seam at one place -- "how a node becomes a tunnel" --
// instead of stubbing out half the program.

#include <asio.hpp>

#include <atomic>
#include <chrono>
#include <cstdio>
#include <fstream>
#include <memory>
#include <string>
#include <vector>

#include "common/config.h"
#include "common/error.h"
#include "egress/direct_egress.h"
#include "egress/egress_manager.h"
#include "harness.h"
#include "selector/history.h"
#include "selector/selector.h"
#include "vpngate/node_store.h"

using namespace ovg;
using namespace std::chrono_literals;

namespace {

using Clock = std::chrono::steady_clock;

bool run_until(asio::io_context &io, const std::function<bool()> &pred,
               std::chrono::milliseconds budget) {
  const auto deadline = Clock::now() + budget;
  while (!pred() && Clock::now() < deadline) {
    io.restart();
    io.run_for(2ms);
  }
  return pred();
}

void settle(asio::io_context &io, std::chrono::milliseconds d = 40ms) {
  const auto deadline = Clock::now() + d;
  while (Clock::now() < deadline) {
    io.restart();
    io.run_for(2ms);
  }
}

// ---------------------------------------------------------------------------
// Loopback peers
// ---------------------------------------------------------------------------

// Accepts forever and echoes. Used both as a target for the direct egress and
// as something the selector's latency prober can actually complete a handshake
// against.
class EchoServer {
 public:
  explicit EchoServer(asio::io_context &io)
      : io_(io), acceptor_(io, asio::ip::tcp::endpoint(
                                   asio::ip::make_address("127.0.0.1"), 0)) {
    accept();
  }

  uint16_t port() const { return acceptor_.local_endpoint().port(); }
  void stop() {
    std::error_code ignored;
    acceptor_.close(ignored);
  }

 private:
  struct Conn : std::enable_shared_from_this<Conn> {
    explicit Conn(asio::ip::tcp::socket s) : sock(std::move(s)) {}
    asio::ip::tcp::socket sock;
    char buf[512];
    void go() {
      auto self = shared_from_this();
      sock.async_read_some(
          asio::buffer(buf), [self](const std::error_code &ec, size_t n) {
            if (ec) return;
            asio::async_write(self->sock, asio::buffer(self->buf, n),
                              [self](const std::error_code &wec, size_t) {
                                if (!wec) self->go();
                              });
          });
    }
  };

  void accept() {
    acceptor_.async_accept([this](const std::error_code &ec,
                                  asio::ip::tcp::socket s) {
      if (ec) return;
      std::make_shared<Conn>(std::move(s))->go();
      accept();
    });
  }

  asio::io_context &io_;
  asio::ip::tcp::acceptor acceptor_;
};

class UdpEcho {
 public:
  explicit UdpEcho(asio::io_context &io)
      : sock_(io, asio::ip::udp::endpoint(asio::ip::make_address("127.0.0.1"),
                                          0)) {
    recv();
  }
  uint16_t port() const { return sock_.local_endpoint().port(); }
  void stop() {
    std::error_code ignored;
    sock_.close(ignored);
  }

 private:
  void recv() {
    sock_.async_receive_from(asio::buffer(buf_), from_,
                             [this](const std::error_code &ec, size_t n) {
                               if (ec) return;
                               std::error_code ignored;
                               sock_.send_to(asio::buffer(buf_, n), from_, 0,
                                             ignored);
                               recv();
                             });
  }
  asio::ip::udp::socket sock_;
  asio::ip::udp::endpoint from_;
  char buf_[2048];
};

// ---------------------------------------------------------------------------
// A synthetic node list whose remotes point at a loopback listener
// ---------------------------------------------------------------------------

std::string base64_encode(const std::string &in) {
  static const char *tbl =
      "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  std::string out;
  size_t i = 0;
  for (; i + 2 < in.size(); i += 3) {
    const uint32_t v = (static_cast<uint8_t>(in[i]) << 16) |
                       (static_cast<uint8_t>(in[i + 1]) << 8) |
                       static_cast<uint8_t>(in[i + 2]);
    out += tbl[(v >> 18) & 63];
    out += tbl[(v >> 12) & 63];
    out += tbl[(v >> 6) & 63];
    out += tbl[v & 63];
  }
  if (i < in.size()) {
    uint32_t v = static_cast<uint8_t>(in[i]) << 16;
    const bool two = i + 1 < in.size();
    if (two) v |= static_cast<uint8_t>(in[i + 1]) << 8;
    out += tbl[(v >> 18) & 63];
    out += tbl[(v >> 12) & 63];
    out += two ? tbl[(v >> 6) & 63] : '=';
    out += '=';
  }
  return out;
}

struct FakeFeed {
  std::string path;
  std::vector<std::string> ids;  // node ids, best prior first

  ~FakeFeed() {
    if (!path.empty()) std::remove(path.c_str());
  }
};

// Three nodes, all of whose OpenVPN profiles say "connect to this loopback
// port". Nothing here ever speaks OpenVPN -- the port only has to complete a
// TCP handshake, which is exactly what the prober measures.
std::string feed_csv(uint16_t port) {
  std::string csv =
      "*vpn_servers\n"
      "#HostName,IP,Score,Ping,Speed,CountryLong,CountryShort,NumVpnSessions,"
      "Uptime,TotalUsers,TotalTraffic,LogType,Operator,Message,"
      "OpenVPN_ConfigData_Base64\n";

  struct Row {
    const char *host;
    const char *ip;
    const char *score;
    const char *country_long;
    const char *country_short;
  };
  const Row rows[] = {
      {"fake-a", "192.0.2.1", "9000000", "Japan", "JP"},
      {"fake-b", "192.0.2.2", "6000000", "Korea", "KR"},
      {"fake-c", "192.0.2.3", "3000000", "Taiwan", "TW"},
  };

  for (const auto &r : rows) {
    const std::string profile =
        std::string("client\ndev tun\nproto tcp\nremote 127.0.0.1 ") +
        std::to_string(port) + " tcp\nnobind\n";
    csv += std::string(r.host) + "," + r.ip + "," + r.score +
           ",10,100000000," + r.country_long + "," + r.country_short +
           ",50,86400000,1000,100000,2weeks,tester,," +
           base64_encode(profile) + "\n";
  }
  csv += "*\n";
  return csv;
}

std::shared_ptr<FakeFeed> write_feed(uint16_t port) {
  auto feed = std::make_shared<FakeFeed>();
  feed->path = "/tmp/ovg_test_feed_" + std::to_string(port) + ".csv";
  feed->ids = {"fake-a@192.0.2.1", "fake-b@192.0.2.2", "fake-c@192.0.2.3"};

  std::ofstream out(feed->path, std::ios::trunc);
  out << feed_csv(port);
  out.close();
  return feed;
}

// Serves the node list over HTTP -- but only once a test says so. Before that
// it answers 503, which is how the store sees a cold start: the process is up,
// the API is reachable, and there is still no directory to select from.
class FeedServer {
 public:
  explicit FeedServer(asio::io_context &io)
      : acceptor_(io, asio::ip::tcp::endpoint(
                          asio::ip::make_address("127.0.0.1"), 0)) {
    acceptor_.listen();
    accept();
  }

  std::string url() const {
    return "http://127.0.0.1:" +
           std::to_string(acceptor_.local_endpoint().port()) + "/api/iphone/";
  }
  void arm(std::string csv) { body_ = std::move(csv); }
  void stop() {
    std::error_code ec;
    acceptor_.close(ec);
  }

 private:
  void accept() {
    auto sock =
        std::make_shared<asio::ip::tcp::socket>(acceptor_.get_executor());
    acceptor_.async_accept(*sock, [this, sock](std::error_code ec) {
      if (ec) return;
      auto resp = std::make_shared<std::string>(
          body_.empty()
              ? std::string("HTTP/1.1 503 Service Unavailable\r\n"
                            "Content-Length: 0\r\n\r\n")
              : "HTTP/1.1 200 OK\r\nContent-Length: " +
                    std::to_string(body_.size()) + "\r\n\r\n" + body_);
      auto buf = std::make_shared<asio::streambuf>();
      asio::async_read_until(
          *sock, *buf, "\r\n\r\n", [sock, resp, buf](std::error_code, size_t) {
            asio::async_write(*sock, asio::buffer(*resp),
                              [sock, resp](std::error_code, size_t) {
                                std::error_code ig;
                                sock->shutdown(
                                    asio::ip::tcp::socket::shutdown_both, ig);
                                sock->close(ig);
                              });
          });
      accept();
    });
  }

  asio::ip::tcp::acceptor acceptor_;
  std::string body_;
};

// ---------------------------------------------------------------------------
// A fake egress: no tunnel, no sockets, fully observable
// ---------------------------------------------------------------------------

class FakeEgress final : public egress::Egress {
 public:
  FakeEgress(std::string node_id, asio::io_context &io)
      : node_id_(std::move(node_id)), io_(io), label_(node_id_) {}

  void async_connect_tcp(const asio::any_io_executor &ex, const Endpoint &,
                         Millis, ConnectHandler h) override {
    asio::post(ex, [h = std::move(h)]() mutable {
      h(make_error_code(Error::NotSupported), nullptr);
    });
  }
  void async_bind_udp(const asio::any_io_executor &ex,
                      UdpBindHandler h) override {
    asio::post(ex, [h = std::move(h)]() mutable {
      h(make_error_code(Error::NotSupported), nullptr);
    });
  }
  void async_resolve(const asio::any_io_executor &ex, const std::string &,
                     ResolveHandler h) override {
    asio::post(ex, [h = std::move(h)]() mutable {
      h(make_error_code(Error::NotSupported), {});
    });
  }

  egress::EgressState state() const override { return state_.load(); }

  egress::EgressStats stats() const override {
    egress::EgressStats s;
    s.node_id = node_id_;
    s.proto = "fake";
    return s;
  }

  std::string detail() const override { return {}; }

  void begin_drain() override {
    auto want = egress::EgressState::Ready;
    state_.compare_exchange_strong(want, egress::EgressState::Draining);
    drained.store(true);
  }

  void shutdown(std::function<void()> on_done) override {
    state_.store(egress::EgressState::Down);
    shutdowns.fetch_add(1);
    if (on_shutdown) on_shutdown(node_id_);
    if (on_done) asio::post(io_, std::move(on_done));
  }

  const std::string &label() const override { return label_; }

  std::atomic<egress::EgressState> state_{egress::EgressState::Ready};
  std::atomic<bool> drained{false};
  std::atomic<int> shutdowns{0};
  // Reported out rather than polled, because the interesting case is the one
  // where the test is deliberately holding no reference to this object.
  std::function<void(const std::string &)> on_shutdown;

 private:
  std::string node_id_;
  asio::io_context &io_;
  std::string label_;
};

// Everything a manager test needs, wired up once.
struct ManagerFixture {
  asio::io_context io;
  EchoServer echo{io};
  // Unused unless a test points cfg.vpngate.api_urls at it; see the cold-start
  // test. Constructed unconditionally because it needs `io`, and `io` is only
  // in scope inside the fixture.
  FeedServer feed_srv{io};
  std::shared_ptr<FakeFeed> feed;
  Config cfg;
  std::unique_ptr<vpngate::NodeStore> store;
  std::unique_ptr<selector::HistoryStore> history;
  std::unique_ptr<selector::Selector> selector;
  std::unique_ptr<egress::EgressManager> mgr;

  // Node ids the factory should refuse, and with which error.
  std::map<std::string, std::error_code> fail_with;
  // Node ids that fail exactly once, then succeed. Keyed by id, value is the
  // remaining number of failures.
  std::map<std::string, int> fail_once;

  // Labels only: holding an egress here would keep the reference count off zero
  // and the drain would never finish.
  std::vector<std::pair<std::string, std::string>> promotions;
  std::vector<std::string> drain_expired;
  std::vector<std::string> shut_down;

  // `tweak` runs after the defaults are in place and before the store is built,
  // which is the only window in which cfg.vpngate still matters: NodeStore
  // copies it and loads the cache in its constructor.
  explicit ManagerFixture(std::function<void(ManagerFixture &)> tweak = {}) {
    feed = write_feed(echo.port());

    cfg.vpngate.api_urls = {"http://127.0.0.1:1/"};
    cfg.vpngate.cache_path = feed->path;
    cfg.vpngate.cache_max_age = std::chrono::hours(24);
    cfg.vpngate.refresh_interval = std::chrono::hours(24);
    cfg.selector.probe_samples = 1;
    cfg.selector.probe_timeout = 500ms;
    cfg.switching.min_interval = 0ms;
    cfg.switching.backoff_initial = 10ms;
    cfg.switching.backoff_max = 40ms;

    if (tweak) tweak(*this);

    store = std::make_unique<vpngate::NodeStore>(io, cfg.vpngate);
    store->start();
    history = std::make_unique<selector::HistoryStore>("", cfg.selector);
    selector = std::make_unique<selector::Selector>(io, cfg.selector, *store,
                                                    *history);
  }

  ~ManagerFixture() {
    if (mgr) {
      bool done = false;
      mgr->shutdown([&] { done = true; });
      run_until(io, [&] { return done; }, 2s);
    }
    store->stop();
    echo.stop();
    feed_srv.stop();
    settle(io);
  }

  // Built after the test has adjusted cfg.
  void build() {
    mgr = std::make_unique<egress::EgressManager>(io, cfg, selector.get(),
                                                  history.get(), nullptr);
    mgr->set_egress_factory(
        [this](const vpngate::Node &node, const vpngate::Remote &,
               egress::EgressManager::NodeReadyHandler on_ready)
            -> egress::EgressPtr {
          const std::string id = node.id();
          auto it = fail_with.find(id);
          if (it != fail_with.end()) {
            auto ec = it->second;
            asio::post(io, [on_ready, ec] { on_ready(ec, "test: refused"); });
            // Non-null: the failure is reported through on_ready, which is the
            // path a real tunnel that comes up and then dies also takes.
            return std::make_shared<FakeEgress>(id, io);
          }
          auto once = fail_once.find(id);
          if (once != fail_once.end() && once->second > 0) {
            once->second--;
            asio::post(io, [on_ready] {
              on_ready(make_error_code(Error::ResourceExhausted),
                       "test: address collision");
            });
            return std::make_shared<FakeEgress>(id, io);
          }
          auto e = std::make_shared<FakeEgress>(id, io);
          e->on_shutdown = [this](const std::string &n) {
            shut_down.push_back(n);
          };
          asio::post(io, [on_ready] { on_ready({}, ""); });
          return e;
        });
    mgr->set_on_promote([this](const egress::EgressPtr &old_e,
                               const egress::EgressPtr &new_e) {
      promotions.emplace_back(old_e ? old_e->label() : "", new_e->label());
    });
    mgr->set_on_drain_expired([this](const egress::EgressPtr &e) {
      drain_expired.push_back(e->label());
    });
  }

  std::error_code start() {
    std::error_code got = make_error_code(Error::Internal);
    bool done = false;
    mgr->start([&](const std::error_code &ec) {
      got = ec;
      done = true;
    });
    run_until(io, [&] { return done; }, 10s);
    return got;
  }
};

}  // namespace

// ---------------------------------------------------------------------------
// DirectEgress
// ---------------------------------------------------------------------------

OVG_TEST(direct_egress_tcp_roundtrip) {
  asio::io_context io;
  EchoServer echo(io);
  auto e = egress::DirectEgress::create(io, DnsConfig{});

  netstack::TcpStreamPtr stream;
  std::error_code connect_ec = make_error_code(Error::Internal);
  bool connected = false;
  e->async_connect_tcp(
      io.get_executor(),
      Endpoint(*IpAddress::parse("127.0.0.1"), echo.port()), 2000ms,
      [&](const std::error_code &ec, netstack::TcpStreamPtr s) {
        connect_ec = ec;
        stream = std::move(s);
        connected = true;
      });
  CHECK(run_until(io, [&] { return connected; }, 3s));
  CHECK_EQ(connect_ec, std::error_code{});
  CHECK(stream != nullptr);
  CHECK(stream->is_open());

  const std::string msg = "hello egress";
  bool wrote = false;
  stream->async_write(asio::buffer(msg),
                      [&](const std::error_code &ec, size_t n) {
                        CHECK_EQ(ec, std::error_code{});
                        CHECK_EQ(n, msg.size());
                        wrote = true;
                      });
  CHECK(run_until(io, [&] { return wrote; }, 3s));

  char rbuf[64];
  size_t got = 0;
  bool read_done = false;
  stream->async_read_some(asio::buffer(rbuf),
                          [&](const std::error_code &ec, size_t n) {
                            CHECK_EQ(ec, std::error_code{});
                            got = n;
                            read_done = true;
                          });
  CHECK(run_until(io, [&] { return read_done; }, 3s));
  CHECK_EQ(std::string(rbuf, got), msg);

  CHECK_EQ(stream->bytes_written(), msg.size());
  CHECK_EQ(stream->bytes_read(), msg.size());

  auto st = e->stats();
  CHECK_EQ(st.tcp_opened, uint64_t{1});
  CHECK_EQ(st.tcp_failed, uint64_t{0});
  CHECK_EQ(st.tcp_active, int64_t{1});

  // The active gauge counts live streams, not live handlers: it only drops when
  // the object is actually destroyed.
  stream.reset();
  settle(io);
  CHECK_EQ(e->stats().tcp_active, int64_t{0});

  echo.stop();
  settle(io);
}

OVG_TEST(direct_egress_half_close_propagates) {
  asio::io_context io;
  EchoServer echo(io);
  auto e = egress::DirectEgress::create(io, DnsConfig{});

  netstack::TcpStreamPtr stream;
  bool connected = false;
  e->async_connect_tcp(io.get_executor(),
                       Endpoint(*IpAddress::parse("127.0.0.1"), echo.port()),
                       2000ms,
                       [&](const std::error_code &ec, netstack::TcpStreamPtr s) {
                         CHECK_EQ(ec, std::error_code{});
                         stream = std::move(s);
                         connected = true;
                       });
  CHECK(run_until(io, [&] { return connected; }, 3s));

  // Our FIN makes the echo server's read fail, so it drops the connection and
  // we see EOF -- the exact sequence a SOCKS5 relay must forward rather than
  // treat as an error.
  stream->shutdown_send();
  char rbuf[16];
  std::error_code read_ec;
  bool read_done = false;
  stream->async_read_some(asio::buffer(rbuf),
                          [&](const std::error_code &ec, size_t) {
                            read_ec = ec;
                            read_done = true;
                          });
  CHECK(run_until(io, [&] { return read_done; }, 3s));
  CHECK_EQ(read_ec, std::error_code(asio::error::eof));

  echo.stop();
  settle(io);
}

OVG_TEST(direct_egress_reports_connect_and_resolve_failures_separately) {
  asio::io_context io;
  auto e = egress::DirectEgress::create(io, DnsConfig{});

  // Port 1 on loopback: nothing listens, and the RST is immediate.
  std::error_code connect_ec;
  bool done = false;
  e->async_connect_tcp(io.get_executor(),
                       Endpoint(*IpAddress::parse("127.0.0.1"), 1), 2000ms,
                       [&](const std::error_code &ec, netstack::TcpStreamPtr s) {
                         connect_ec = ec;
                         CHECK(s == nullptr);
                         done = true;
                       });
  CHECK(run_until(io, [&] { return done; }, 5s));
  CHECK_NE(connect_ec, std::error_code{});
  CHECK_NE(connect_ec, std::error_code(make_error_code(Error::ResolveFailed)));
  CHECK_EQ(e->stats().tcp_failed, uint64_t{1});
  CHECK_EQ(e->stats().tcp_opened, uint64_t{0});

  // A name that cannot resolve must not be reported as a connect failure: the
  // SOCKS5 reply code and the operator's debugging both depend on the
  // difference. Some environments answer for .invalid anyway (a wildcard
  // resolver, a captive DNS), and then there is nothing here to assert.
  std::error_code resolve_ec;
  bool resolved_anyway = false;
  bool done2 = false;
  e->async_resolve(io.get_executor(), "no-such-host.invalid",
                   [&](const std::error_code &ec, std::vector<IpAddress> a) {
                     resolve_ec = ec;
                     resolved_anyway = !ec && !a.empty();
                     done2 = true;
                   });
  if (!run_until(io, [&] { return done2; }, 10s)) {
    SKIP("resolver did not answer for .invalid within 10s");
  }
  if (resolved_anyway) SKIP("this environment resolves .invalid");
  CHECK_EQ(resolve_ec, std::error_code(make_error_code(Error::ResolveFailed)));

  std::error_code by_name_ec;
  bool done3 = false;
  e->async_connect_tcp(io.get_executor(), Endpoint("no-such-host.invalid", 80),
                       2000ms,
                       [&](const std::error_code &ec, netstack::TcpStreamPtr s) {
                         by_name_ec = ec;
                         CHECK(s == nullptr);
                         done3 = true;
                       });
  CHECK(run_until(io, [&] { return done3; }, 10s));
  CHECK_EQ(by_name_ec, std::error_code(make_error_code(Error::ResolveFailed)));
}

OVG_TEST(direct_egress_udp_roundtrip) {
  asio::io_context io;
  UdpEcho peer(io);
  auto e = egress::DirectEgress::create(io, DnsConfig{});

  netstack::UdpSocketPtr sock;
  bool bound = false;
  e->async_bind_udp(io.get_executor(),
                    [&](const std::error_code &ec, netstack::UdpSocketPtr s) {
                      CHECK_EQ(ec, std::error_code{});
                      sock = std::move(s);
                      bound = true;
                    });
  CHECK(run_until(io, [&] { return bound; }, 3s));
  CHECK(sock != nullptr);
  CHECK_EQ(e->stats().udp_active, int64_t{1});

  char rbuf[64];
  size_t got = 0;
  Endpoint from;
  bool received = false;
  sock->async_receive_from(asio::buffer(rbuf),
                           [&](const std::error_code &ec, size_t n,
                               const Endpoint &f) {
                             CHECK_EQ(ec, std::error_code{});
                             got = n;
                             from = f;
                             received = true;
                           });

  const std::string msg = "datagram";
  bool sent = false;
  sock->async_send_to(asio::buffer(msg),
                      Endpoint(*IpAddress::parse("127.0.0.1"), peer.port()),
                      [&](const std::error_code &ec, size_t n) {
                        CHECK_EQ(ec, std::error_code{});
                        CHECK_EQ(n, msg.size());
                        sent = true;
                      });
  CHECK(run_until(io, [&] { return sent && received; }, 3s));
  CHECK_EQ(std::string(rbuf, got), msg);
  CHECK_EQ(from.port(), peer.port());

  sock.reset();
  settle(io);
  CHECK_EQ(e->stats().udp_active, int64_t{0});
  peer.stop();
  settle(io);
}

OVG_TEST(direct_egress_draining_refuses_new_work) {
  asio::io_context io;
  EchoServer echo(io);
  auto e = egress::DirectEgress::create(io, DnsConfig{});

  e->begin_drain();
  CHECK_EQ(e->state(), egress::EgressState::Draining);
  CHECK(!e->usable());

  // Draining is not "gone": the distinction is what tells the SOCKS5 layer to
  // retry on the new egress rather than report a failure to the client.
  std::error_code ec1;
  bool done1 = false;
  e->async_connect_tcp(io.get_executor(),
                       Endpoint(*IpAddress::parse("127.0.0.1"), echo.port()),
                       1000ms,
                       [&](const std::error_code &ec, netstack::TcpStreamPtr) {
                         ec1 = ec;
                         done1 = true;
                       });
  CHECK(run_until(io, [&] { return done1; }, 2s));
  CHECK_EQ(ec1, std::error_code(make_error_code(Error::EgressDraining)));

  bool shut = false;
  e->shutdown([&] { shut = true; });
  CHECK(run_until(io, [&] { return shut; }, 2s));
  CHECK_EQ(e->state(), egress::EgressState::Down);

  std::error_code ec2;
  bool done2 = false;
  e->async_bind_udp(io.get_executor(),
                    [&](const std::error_code &ec, netstack::UdpSocketPtr) {
                      ec2 = ec;
                      done2 = true;
                    });
  CHECK(run_until(io, [&] { return done2; }, 2s));
  CHECK_EQ(ec2, std::error_code(make_error_code(Error::EgressGone)));

  echo.stop();
  settle(io);
}

// ---------------------------------------------------------------------------
// EgressManager: direct mode
// ---------------------------------------------------------------------------

OVG_TEST(manager_direct_mode_needs_no_selector) {
  asio::io_context io;
  Config cfg;
  cfg.egress_mode = "direct";
  egress::EgressManager mgr(io, cfg, nullptr, nullptr, nullptr);

  std::error_code start_ec = make_error_code(Error::Internal);
  bool started = false;
  mgr.start([&](const std::error_code &ec) {
    start_ec = ec;
    started = true;
  });
  CHECK(run_until(io, [&] { return started; }, 3s));
  CHECK_EQ(start_ec, std::error_code{});

  auto e = mgr.acquire();
  CHECK(e != nullptr);
  CHECK(e->usable());
  CHECK_EQ(e->label(), std::string("direct"));

  // There is nothing to switch to, and pretending otherwise would mean
  // silently dropping sessions for no gain.
  CHECK(!mgr.request_switch(egress::SwitchReason::Manual, /*force=*/true));

  bool shut = false;
  mgr.shutdown([&] { shut = true; });
  CHECK(run_until(io, [&] { return shut; }, 3s));
  CHECK(mgr.acquire() == nullptr);
}

// ---------------------------------------------------------------------------
// EgressManager: make-before-break
// ---------------------------------------------------------------------------

OVG_TEST(manager_promotes_without_touching_existing_sessions) {
  ManagerFixture fx;
  fx.cfg.switching.drain_grace = 30s;  // long enough that nothing expires here
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto first = fx.mgr->acquire();
  CHECK(first != nullptr);
  CHECK(first->usable());

  // This is a session: it holds the egress for its whole life, and that
  // reference is the only thing keeping the old tunnel alive after the switch.
  auto session = fx.mgr->acquire();
  CHECK_EQ(session.get(), first.get());

  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual));
  CHECK(run_until(
      fx.io, [&] { return fx.mgr->acquire() != nullptr &&
                          fx.mgr->acquire().get() != first.get(); },
      10s));

  auto second = fx.mgr->acquire();
  CHECK(second != nullptr);
  CHECK_NE(second->label(), first->label());

  // The old egress refuses new work but was not torn down: the session on it is
  // still connected, which is the entire point of make-before-break.
  CHECK_EQ(first->state(), egress::EgressState::Draining);
  CHECK_EQ(std::static_pointer_cast<FakeEgress>(first)->shutdowns.load(), 0);
  CHECK_EQ(fx.mgr->status().draining, size_t{1});

  // The callback fired before the drain started, with both egresses live, so
  // the SOCKS5 layer can re-home what is safe to re-home.
  CHECK_EQ(fx.promotions.size(), size_t{2});  // startup + this switch
  CHECK_EQ(fx.promotions.back().first, first->label());
  CHECK_EQ(fx.promotions.back().second, second->label());

  // Session ends. use_count() falling to the manager's own reference is what
  // completes the drain -- there is no session table to consult. Which is also
  // why the test must let go of every reference of its own first.
  const std::string old_label = first->label();
  std::weak_ptr<egress::Egress> weak_old = first;
  session.reset();
  first.reset();
  CHECK(run_until(fx.io, [&] { return !fx.shut_down.empty(); }, 5s));
  CHECK_EQ(fx.shut_down.back(), old_label);
  CHECK_EQ(fx.mgr->status().draining, size_t{0});
  CHECK(fx.drain_expired.empty());  // it finished, it did not time out
  // And nothing kept it alive: a leak here would be a tunnel that never closes.
  CHECK(weak_old.expired());
}

OVG_TEST(manager_closes_stragglers_when_the_grace_window_expires) {
  ManagerFixture fx;
  fx.cfg.switching.drain_grace = 150ms;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto first = fx.mgr->acquire();
  CHECK(first != nullptr);
  auto stuck_session = first;  // never released: a long-lived download

  CHECK(fx.mgr->request_switch(egress::SwitchReason::Unhealthy));
  CHECK(run_until(
      fx.io, [&] { auto a = fx.mgr->acquire();
                   return a && a.get() != first.get(); }, 10s));

  // The requirement's own fallback: sessions that cannot be carried across are
  // dropped rather than kept alive on a node we have decided is bad.
  CHECK(run_until(fx.io, [&] { return !fx.drain_expired.empty(); }, 5s));
  CHECK_EQ(fx.drain_expired.front(), first->label());
  CHECK_EQ(std::static_pointer_cast<FakeEgress>(first)->shutdowns.load(), 1);
  CHECK_EQ(first->state(), egress::EgressState::Down);
  CHECK_EQ(fx.mgr->status().draining, size_t{0});
  (void)stuck_session;
}

OVG_TEST(manager_hard_mode_closes_everything_at_promotion) {
  ManagerFixture fx;
  fx.cfg.switching.mode = SwitchConfig::Mode::Hard;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto first = fx.mgr->acquire();
  CHECK(first != nullptr);
  auto session = first;

  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual));
  CHECK(run_until(
      fx.io, [&] { auto a = fx.mgr->acquire();
                   return a && a.get() != first.get(); }, 10s));

  // No drain at all: the old egress is down before the switch is reported, and
  // the SOCKS5 layer was told to close what was on it.
  CHECK_EQ(first->state(), egress::EgressState::Down);
  CHECK_EQ(fx.mgr->status().draining, size_t{0});
  CHECK_EQ(fx.drain_expired.size(), size_t{1});
  CHECK_EQ(fx.drain_expired.front(), first->label());
  (void)session;
}

OVG_TEST(manager_walks_the_candidate_list_when_one_fails) {
  ManagerFixture fx;
  // Refuse the two highest-scoring nodes; the third must still bring the
  // service up. A failure here costs nothing, because nothing was promoted.
  fx.fail_with[fx.feed->ids[0]] = make_error_code(Error::TunnelSetupFailed);
  fx.fail_with[fx.feed->ids[1]] = make_error_code(Error::Timeout);
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto e = fx.mgr->acquire();
  CHECK(e != nullptr);
  CHECK_EQ(e->label(), fx.feed->ids[2]);
  CHECK_EQ(fx.mgr->status().switch_failures, uint64_t{0});
}

OVG_TEST(manager_gives_up_startup_after_the_retry_budget) {
  ManagerFixture fx;
  for (const auto &id : fx.feed->ids)
    fx.fail_with[id] = make_error_code(Error::TunnelSetupFailed);
  fx.build();

  // Reporting the failure rather than retrying forever: a service that never
  // finishes starting is harder to operate than one that exits with a reason.
  const auto ec = fx.start();
  CHECK_EQ(ec, std::error_code(make_error_code(Error::TunnelSetupFailed)));
  CHECK(fx.mgr->acquire() == nullptr);
  CHECK_GT(fx.mgr->status().switch_failures, uint64_t{0});
}

OVG_TEST(manager_waits_for_the_directory_instead_of_spending_startup_rounds) {
  // Regression: the first selection normally loses the race with the first
  // directory fetch by about a second. Charging that to the backoff ladder cost
  // a full backoff -- and four of them exhausted the startup budget, so a
  // process whose API call took a few seconds gave up and exited with
  // "no node came up after 4 rounds" while the node list was still in flight.
  // Nothing had been attempted. It is a wait, not a failure.
  std::string cache;
  ManagerFixture fx([&](ManagerFixture &f) {
    // No cache on disk and nothing served yet: a genuinely cold start.
    cache = f.feed->path + ".cold";
    std::remove(cache.c_str());
    f.cfg.vpngate.cache_path = cache;
    f.cfg.vpngate.api_urls = {f.feed_srv.url()};
    // Small enough that the buggy path would have burned all four rounds
    // (5 + 10 + 10 + 10 ms) long before the check below.
    f.cfg.switching.backoff_initial = 5ms;
    f.cfg.switching.backoff_max = 10ms;
  });
  fx.build();

  std::error_code got = make_error_code(Error::Internal);
  bool done = false;
  fx.mgr->start([&](const std::error_code &ec) {
    got = ec;
    done = true;
  });

  // Still starting, with nothing charged against it.
  CHECK(!run_until(fx.io, [&] { return done; }, 200ms));
  CHECK_EQ(fx.mgr->status().switch_failures, uint64_t{0});
  // And reachable: the wait releases the phase, so the tunnel-down watchdog and
  // an admin switch are not told a selection is in progress for the whole wait.
  CHECK(fx.mgr->status().phase == egress::SwitchPhase::Idle);

  fx.feed_srv.arm(feed_csv(fx.echo.port()));
  fx.store->refresh_now(nullptr);

  // The poll is on a one-second cadence, so allow a couple of them.
  CHECK(run_until(fx.io, [&] { return done; }, 4s));
  CHECK_EQ(got, std::error_code{});
  CHECK_EQ(fx.promotions.size(), size_t{1});
  CHECK_EQ(fx.mgr->status().switch_failures, uint64_t{0});
  CHECK(fx.mgr->acquire() != nullptr);

  std::remove(cache.c_str());
}

OVG_TEST(manager_address_collision_degrades_to_a_hard_switch) {
  ManagerFixture fx;
  fx.cfg.switching.drain_grace = 30s;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto first = fx.mgr->acquire();
  CHECK(first != nullptr);
  auto session = first;

  // Every remaining node reports the collision once. lwIP cannot tell two
  // tunnels pushing the same private address apart (netstack/lwip_stack.h), so
  // make-before-break is simply unavailable and the switch must degrade rather
  // than stall.
  for (const auto &id : fx.feed->ids) fx.fail_once[id] = 1;

  CHECK(fx.mgr->request_switch(egress::SwitchReason::Unhealthy));
  CHECK(run_until(
      fx.io, [&] { auto a = fx.mgr->acquire();
                   return a && a.get() != first.get(); }, 15s));

  // The old egress was closed to make room, and the layer above was told so it
  // could close the sessions that were on it.
  CHECK_EQ(first->state(), egress::EgressState::Down);
  CHECK_GT(fx.drain_expired.size(), size_t{0});
  CHECK_EQ(fx.drain_expired.front(), first->label());
  CHECK(fx.mgr->acquire()->usable());
  (void)session;
}

OVG_TEST(manager_bounds_the_number_of_draining_egresses) {
  ManagerFixture fx;
  fx.cfg.switching.drain_grace = 30s;
  fx.cfg.switching.max_draining = 1;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto a = fx.mgr->acquire();
  auto session_a = a;  // held: a would otherwise drain immediately
  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual));
  CHECK(run_until(fx.io, [&] { auto x = fx.mgr->acquire();
                               return x && x.get() != a.get(); }, 10s));

  auto b = fx.mgr->acquire();
  auto session_b = b;
  CHECK_EQ(fx.mgr->status().draining, size_t{1});

  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual));
  CHECK(run_until(fx.io, [&] { auto x = fx.mgr->acquire();
                               return x && x.get() != b.get(); }, 10s));

  // Each draining tunnel still costs an OpenVPN session, a netif and its
  // buffers; without a cap a flapping node would accumulate them until memory
  // ran out. The oldest goes, even though its grace window had not expired.
  CHECK_EQ(fx.mgr->status().draining, size_t{1});
  CHECK_EQ(a->state(), egress::EgressState::Down);
  CHECK_EQ(b->state(), egress::EgressState::Draining);
  (void)session_a;
  (void)session_b;
}

OVG_TEST(manager_suppresses_flapping_but_honours_force) {
  ManagerFixture fx;
  fx.cfg.switching.min_interval = 60s;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});
  CHECK(fx.mgr->acquire() != nullptr);

  // A health monitor that fires every 15s must not be able to switch every 15s.
  CHECK(!fx.mgr->request_switch(egress::SwitchReason::Unhealthy));
  CHECK(!fx.mgr->request_switch(egress::SwitchReason::BetterCandidate));

  // The admin endpoint can override it, because an operator asking for a switch
  // knows something the anti-flap timer does not.
  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual, /*force=*/true));
  auto first = fx.mgr->acquire();
  CHECK(run_until(fx.io, [&] { auto x = fx.mgr->acquire();
                               return x && x.get() != first.get(); }, 10s));
}

OVG_TEST(manager_opportunistic_switch_requires_a_better_score) {
  ManagerFixture fx;
  // A challenger has to beat the incumbent by a wide margin. Since all three
  // synthetic nodes probe the same loopback listener, none can, so an
  // opportunistic request must decide to stay put -- and must not be recorded
  // as a failure, because nothing failed.
  fx.cfg.switching.improvement_margin = 5.0;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto before = fx.mgr->acquire();
  CHECK(before != nullptr);
  CHECK(fx.mgr->request_switch(egress::SwitchReason::BetterCandidate));
  settle(fx.io, 400ms);

  CHECK_EQ(fx.mgr->acquire().get(), before.get());
  CHECK_EQ(fx.mgr->status().switch_failures, uint64_t{0});
  CHECK_EQ(fx.mgr->status().phase, egress::SwitchPhase::Idle);

  // Unhealthy is not opportunistic: when the current node is bad, anything that
  // connects is an improvement and the margin does not apply.
  CHECK(fx.mgr->request_switch(egress::SwitchReason::Unhealthy));
  CHECK(run_until(fx.io, [&] { auto x = fx.mgr->acquire();
                               return x && x.get() != before.get(); }, 10s));
}

OVG_TEST(manager_never_hands_out_a_draining_egress) {
  ManagerFixture fx;
  fx.cfg.switching.drain_grace = 30s;
  fx.build();
  CHECK_EQ(fx.start(), std::error_code{});

  auto first = fx.mgr->acquire();
  auto session = first;
  CHECK(fx.mgr->request_switch(egress::SwitchReason::Manual));
  CHECK(run_until(fx.io, [&] { auto x = fx.mgr->acquire();
                               return x && x.get() != first.get(); }, 10s));

  // A hundred new sessions all land on the new egress. If any of them got the
  // draining one, its drain would never finish.
  for (int i = 0; i < 100; ++i) {
    auto e = fx.mgr->acquire();
    CHECK(e != nullptr);
    CHECK_NE(e.get(), first.get());
    CHECK(e->usable());
  }
  (void)session;
}