#include <doctest/doctest.h>

#include "packeteer/l7/dissector.hpp"

using namespace packeteer::net;

namespace {

// A dissector that claims a port but never actually matches any
// payload - stands in for e.g. TlsSniDissector receiving QUIC bytes
// on port 443: same port, wrong protocol, never succeeds.
class NeverMatchesDissector : public L7Dissector {
public:
    explicit NeverMatchesDissector(std::uint16_t port) : port_(port) {}
    std::uint16_t port() const override { return port_; }
    std::optional<std::string> summarize(std::span<const unsigned char>) const override {
        return std::nullopt;
    }

private:
    std::uint16_t port_;
};

class AlwaysMatchesDissector : public L7Dissector {
public:
    AlwaysMatchesDissector(std::uint16_t port, std::string label, Transport transport = Transport::kAny)
        : port_(port), label_(std::move(label)), transport_(transport) {}
    std::uint16_t port() const override { return port_; }
    Transport transport() const override { return transport_; }
    std::optional<std::string> summarize(std::span<const unsigned char>) const override {
        return label_;
    }

private:
    std::uint16_t port_;
    std::string label_;
    Transport transport_;
};

}  // namespace

TEST_CASE("L7Registry falls through to a later dissector on the same port "
          "when an earlier one fails to match") {
    NeverMatchesDissector tls_like(443);
    AlwaysMatchesDissector quic_like(443, "QUIC something");

    L7Registry registry;
    registry.add(&tls_like);
    registry.add(&quic_like);

    auto result = registry.dissect(443, Transport::kTcp, {});
    REQUIRE(result.has_value());
    CHECK(*result == "QUIC something");
}

TEST_CASE("L7Registry still returns the first dissector to succeed, not the last") {
    AlwaysMatchesDissector first(80, "first");
    AlwaysMatchesDissector second(80, "second");

    L7Registry registry;
    registry.add(&first);
    registry.add(&second);

    auto result = registry.dissect(80, Transport::kTcp, {});
    REQUIRE(result.has_value());
    CHECK(*result == "first");
}

TEST_CASE("L7Registry returns nullopt when no dissector on the port matches") {
    NeverMatchesDissector only(443);

    L7Registry registry;
    registry.add(&only);

    CHECK_FALSE(registry.dissect(443, Transport::kTcp, {}).has_value());
}

TEST_CASE("L7Registry skips a dissector whose declared transport doesn't match, even on the "
          "right port") {
    // The actual fix for QUIC's false positives on TCP:443: a
    // dissector that only claims UDP must never be tried against a
    // TCP payload on the same port, regardless of what its own
    // summarize() would have returned.
    AlwaysMatchesDissector udp_only(443, "UDP thing", Transport::kUdp);

    L7Registry registry;
    registry.add(&udp_only);

    CHECK_FALSE(registry.dissect(443, Transport::kTcp, {}).has_value());
    CHECK(registry.dissect(443, Transport::kUdp, {}).has_value());
}

TEST_CASE("L7Registry's default kAny transport matches either TCP or UDP") {
    AlwaysMatchesDissector any(53, "DNS-like");  // default transport: kAny

    L7Registry registry;
    registry.add(&any);

    CHECK(registry.dissect(53, Transport::kTcp, {}).has_value());
    CHECK(registry.dissect(53, Transport::kUdp, {}).has_value());
}