#include <balancer/kernel/net/ip_matcher.h>

#include <library/cpp/testing/unittest/registar.h>

Y_UNIT_TEST_SUITE(TIpMatcherTest) {
    using namespace NSrvKernel;

    auto GetIpv4(const TString& host) {
        struct sockaddr_in sin;
        Zero(sin);
        sin.sin_family = AF_INET;
        sin.sin_port = 0;
        UNIT_ASSERT_EQUAL(inet_pton(AF_INET, host.c_str(), &sin.sin_addr), 1);

        return NAddr::TIPv4Addr(sin);
    }

    auto GetIpv6(const TString& host) {
        struct sockaddr_in6 sin6;
        Zero(sin6);
        sin6.sin6_family = AF_INET6;
        sin6.sin6_port = 0;
        UNIT_ASSERT_EQUAL(inet_pton(AF_INET6, host.c_str(), &sin6.sin6_addr), 1);

        return NAddr::TIPv6Addr(sin6);
    }

    Y_UNIT_TEST(TestIpMatcher) {
        auto check = [] (auto&& net) {
            TIpMatcher matcher(net);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("1.1.1.1")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("1.1.1.2")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("2.2.0.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("2.2.255.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("2.3.0.0")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("2.1.255.255")), false);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1:1::1")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1:1::2")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2:1:ffff:ffff:ffff:ffff:ffff:ffff")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2:2::0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2:2:0:0:ffff:ffff:ffff:ffff")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2:2:3::0")), false);
        };

        TString cidrNets = "1.1.1.1,2.2.0.0/16,1:1::1,2:2::/64";
        TString rangeNets = "1.1.1.1-1.1.1.1,2.2.0.0-2.2.255.255,1:1::1-1:1::1,"
                            "2:2::-2:2:0:0:ffff:ffff:ffff:ffff";
        TString mixedNets = "1.1.1.1,2.2.0.0-2.2.255.255,1:1::1-1:1::1,"
                            "2:2::/64";
        check(cidrNets);
        check(rangeNets);
        check(mixedNets);
    }
    
    Y_UNIT_TEST(TestIpMatcherOverflow) {
        auto check = [] (auto&& net) {
            TIpMatcher matcher(net);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("0.0.0.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("1.1.1.1")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("255.255.255.255")), true);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("::")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::1")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")), true);
        };

        TString cidrNets = "0.0.0.0/0,::/0";
        TString rangeNets = "0.0.0.0-255.255.255.255,::-ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff";
        TString mixedNets = "0.0.0.0-255.255.255.255,::/0";
        check(cidrNets);
        check(rangeNets);
        check(mixedNets);
    }

    Y_UNIT_TEST(TestIpMatcherMultipleNetsv4) {
        auto check = [] (auto&& net) {
            TIpMatcher matcher(net);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.191.255")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.192.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.193.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.199.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.200.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.201.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.201.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.202.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.202.100")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.202.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.202.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.0")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.63")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.64")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.100")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.127")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("5.45.203.128")), false);


            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("95.108.181.255")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("95.108.182.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("95.108.182.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("95.108.183.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("95.108.184.0")), false);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("172.23.255.255")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("172.24.0.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("172.26.0.0")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("172.31.255.255")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv4("172.32.0.0")), false);
        };

        TString cidrNets = "5.45.192.0/21,5.45.200.0/23,5.45.202.0/24,"
                           "5.45.203.64/26,95.108.182.0/23,172.24.0.0/13";
        TString rangeNets = "5.45.192.0-5.45.199.255,5.45.200.0-5.45.201.255,"
                            "5.45.202.0-5.45.202.255,5.45.203.64-5.45.203.127,"
                            "95.108.182.0-95.108.183.255,172.24.0.0-172.31.255.255";
        TString mixedNets = "5.45.192.0/21,5.45.200.0-5.45.201.255,5.45.202.0/24,"
                            "5.45.203.64-5.45.203.127,95.108.182.0/23,172.24.0.0-172.31.255.255";
        check(cidrNets);
        check(rangeNets);
        check(mixedNets);
    }
    
    Y_UNIT_TEST(TestIpMatcherMultipleNetsv6) {
        auto check = [] (auto&& net) {
            TIpMatcher matcher(net);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::eff")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::f00")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::f0a")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::fff")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("1::1000")), false);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2::eff")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2::f00")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2::f0a")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2::fff")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("2::1000")), false);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("3::eff")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("3::f00")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("3::f0a")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("3::fff")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("3::1000")), false);

            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("4::eff")), false);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("4::f00")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("4::f0a")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("4::fff")), true);
            UNIT_ASSERT_EQUAL(matcher.Match(GetIpv6("4::1000")), false);
        };

        TString cidrNets = "1::f00/120,2::f00/120,3::f00/120,4::f00/120";
        TString rangeNets = "1::f00-1::fff,2::f00-2::fff,3::f00-3::fff,4::f00-4::fff";
        TString mixedNets = "1::f00-1::fff,2::f00/120,3::f00-3::fff,4::f00/120";
        check(cidrNets);
        check(rangeNets);
        check(mixedNets);
    }
};
