Skip to content

Commit

Permalink
Last functions and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@mohsaka
mohsaka committed Jun 28, 2024
1 parent 305de00 commit 21a8d16
Show file tree
Hide file tree
Showing 2 changed files with 243 additions and 23 deletions.
79 changes: 70 additions & 9 deletions velox/functions/prestosql/IPAddressFunctions.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,13 @@ struct IPAddressFunction {
boost::uuids::random_generator generator_;
};

inline bool isIPV4(int128_t ip){
int128_t ipV4 = 0x0000FFFF00000000;
int128_t mask = 0xFFFFFFFFFFFFFFFF;
mask = (mask << 64) | 0xFFFFFFFF00000000;
return (ip & mask) == ipV4;
}

template <typename T>
struct IPPrefixFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);
Expand Down Expand Up @@ -85,6 +92,29 @@ struct IPPrefixFunction {

result = std::make_shared<IPPrefix>(canonicalAddrInt, prefix);
}

FOLLY_ALWAYS_INLINE void call(
out_type<TheIPPrefix>& result,
const arg_type<Varchar>& ip,
const arg_type<int8_t> prefixBits) {

boost::asio::ip::address_v6::bytes_type addrBytes;
auto addr = boost::asio::ip::make_address(ip);
int128_t intAddr;
if (addr.is_v4()) {
addrBytes = boost::asio::ip::make_address_v6(
boost::asio::ip::v4_mapped, addr.to_v4())
.to_bytes();
} else {
addrBytes = addr.to_v6().to_bytes();
}

bigEndianByteArray(addrBytes);
memcpy(&intAddr, &addrBytes, 16);

call(result, intAddr, prefixBits);

}
};

template <typename T>
Expand All @@ -106,19 +136,11 @@ inline int128_t getIPSubnetMax(int128_t ip, uint8_t prefix){
boost::asio::ip::address_v6::bytes_type addrBytes;
memcpy(&result, &ip, 16);

memcpy(&addrBytes, &ip, 16);
bigEndianByteArray(addrBytes);

auto v6Addr = boost::asio::ip::make_address_v6(addrBytes);

if (v6Addr.is_v4_mapped()) {
assert(prefix <= 32);
if (isIPV4(ip)) {
if (prefix < 32) {
result |= (mask << 32 - prefix) - 1;
}
} else {
assert(prefix <= 128);

// Special case. Return all bits set to 1;
if (prefix == 0) {
result = -1;
Expand Down Expand Up @@ -154,20 +176,59 @@ struct IPSubnetRangeFunction {
}
};

template <typename T>
struct IPSubnetOfFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);
FOLLY_ALWAYS_INLINE void call(
out_type<bool>& result,
const arg_type<TheIPPrefix>& ipPrefix,
const arg_type<IPAddress>& ip) {

uint128_t mask = 1;
uint8_t prefix = (uint8_t)ipPrefix->prefix;
int128_t checkIP = ip;

if (isIPV4(ipPrefix->ip)) {
if (prefix < 32) {
checkIP &= ((mask << 32 - prefix) - 1) ^ -1;
}
} else {
if (prefix < 128) {
checkIP &= ((mask << 128 - prefix) - 1) ^ -1;
}
}
result = (ipPrefix->ip == checkIP);
}

FOLLY_ALWAYS_INLINE void call(
out_type<bool>& result,
const arg_type<TheIPPrefix>& ipPrefix,
const arg_type<TheIPPrefix>& ipPrefix2) {

call(result, ipPrefix, ipPrefix2->ip);
result = result && (ipPrefix2->prefix >= ipPrefix->prefix);

}
};

void registerIPAddressFunctions(const std::string& prefix) {
registerIPAddressType();
registerIPPrefixType();
registerFunction<IPAddressFunction, IPAddress>({prefix + "ipaddress"});
registerFunction<IPPrefixFunction, TheIPPrefix, IPAddress, int8_t>(
{prefix + "ip_prefix"});
registerFunction<IPPrefixFunction, TheIPPrefix, Varchar, int8_t>(
{prefix + "ip_prefix"});
registerFunction<IPSubnetMinFunction, IPAddress, TheIPPrefix>(
{prefix + "ip_subnet_min"});
registerFunction<IPSubnetMaxFunction, IPAddress, TheIPPrefix>(
{prefix + "ip_subnet_max"});
registerFunction<IPSubnetRangeFunction, Array<IPAddress>, TheIPPrefix>(
{prefix + "ip_subnet_range"});
registerFunction<IPSubnetOfFunction, bool, TheIPPrefix, IPAddress>(
{prefix + "is_subnet_of"});
registerFunction<IPSubnetOfFunction, bool, TheIPPrefix, TheIPPrefix>(
{prefix + "is_subnet_of"});
// registerFunction<IPPrefixFunction, TheIPPrefix, Varchar, int8_t>(
// {prefix + "ip_prefix"});
}
Expand Down
187 changes: 173 additions & 14 deletions velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,23 @@ class IPAddressTest : public functions::test::FunctionBaseTest {
"cast(ip_subnet_range(cast(c0 as ipprefix))[2] as varchar)", input);
return result;
}

std::optional<bool> getIsSubnetOfIP(
const std::optional<std::string> prefix,
const std::optional<std::string> ip) {
auto result = evaluateOnce<bool>(
"is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipaddress))", prefix, ip);
return result;
}

std::optional<bool> getIsSubnetOfIPPrefix(
const std::optional<std::string> prefix,
const std::optional<std::string> prefix2) {
auto result = evaluateOnce<bool>(
"is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipprefix))", prefix, prefix2);
return result;
}

};

TEST_F(IPAddressTest, castAsVarchar) {
Expand Down Expand Up @@ -113,11 +130,20 @@ TEST_F(IPAddressTest, IPPrefixv4) {
EXPECT_THROW(getIPPrefix("10.135.23.12", 33), VeloxUserError);
}

TEST_F(IPAddressTest, IPPrefixv4UsingVarchar) {
EXPECT_EQ("10.0.0.0/8", getIPPrefixUsingVarchar("10.135.23.12", 8));
EXPECT_EQ("192.128.0.0/9", getIPPrefixUsingVarchar("192.168.255.255", 9));
EXPECT_EQ("192.168.255.255/32", getIPPrefixUsingVarchar("192.168.255.255", 32));
EXPECT_EQ("0.0.0.0/0", getIPPrefixUsingVarchar("192.168.255.255", 0));

EXPECT_THROW(getIPPrefixUsingVarchar("12.483.09.1", 8), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23.12.12", 8), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23", 8), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("12.135.23.12", -1), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23.12", 33), VeloxUserError);
}

TEST_F(IPAddressTest, IPPrefixv6) {
// EXPECT_EQ(
// "2001:db8:85a3::/48",
// getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334",
// 48));
EXPECT_EQ(
"2001:db8:85a3::/48",
getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 48));
Expand Down Expand Up @@ -145,6 +171,69 @@ TEST_F(IPAddressTest, IPPrefixv6) {
VeloxUserError);
}

TEST_F(IPAddressTest, IPPrefixv6UsingVarchar) {
EXPECT_EQ(
"2001:db8:85a3::/48",
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 48));
EXPECT_EQ(
"2001:db8:85a3::/52",
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 52));
EXPECT_EQ(
"2001:db8:85a3:1:1:8a2e:370:7334/128",
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 128));
EXPECT_EQ("::/0", getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 0));

EXPECT_THROW(
getIPPrefixUsingVarchar("q001:0db8:85a3:0001:0001:8a2e:0370:7334", 8),
VeloxUserError);
EXPECT_THROW(
getIPPrefixUsingVarchar("2001:0db8:85a3:542e:0001:0001:8a2e:0370:7334", 8),
VeloxUserError);
EXPECT_THROW(
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370", 8), VeloxUserError);
EXPECT_THROW(
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", -1),
VeloxUserError);
EXPECT_THROW(
getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 140),
VeloxUserError);
}

TEST_F(IPAddressTest, IPPrefixPrestoTests) {
EXPECT_EQ(getIPPrefix("1.2.3.4", 24), "1.2.3.0/24");
EXPECT_EQ(getIPPrefix( "1.2.3.4", 32), "1.2.3.4/32");
EXPECT_EQ(getIPPrefix( "1.2.3.4", 0), "0.0.0.0/0");
EXPECT_EQ(getIPPrefix( "::ffff:1.2.3.4", 24), "1.2.3.0/24");
EXPECT_EQ(getIPPrefix( "64:ff9b::17", 64), "64:ff9b::/64");
EXPECT_EQ(getIPPrefix( "64:ff9b::17", 127), "64:ff9b::16/127");
EXPECT_EQ(getIPPrefix( "64:ff9b::17", 128), "64:ff9b::17/128");
EXPECT_EQ(getIPPrefix( "64:ff9b::17", 0), "::/0");
EXPECT_THROW(getIPPrefix( "::ffff:1.2.3.4", -1), VeloxUserError);
EXPECT_THROW(getIPPrefix( "::ffff:1.2.3.4", 33), VeloxUserError);
EXPECT_THROW(getIPPrefix( "64:ff9b::10", -1), VeloxUserError);
EXPECT_THROW(getIPPrefix( "64:ff9b::10", 129), VeloxUserError);
}

TEST_F(IPAddressTest, IPPrefixVarcharPrestoTests) {
EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 24), "1.2.3.0/24");
EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 32), "1.2.3.4/32");
EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 0), "0.0.0.0/0");
EXPECT_EQ(getIPPrefixUsingVarchar("::ffff:1.2.3.4", 24), "1.2.3.0/24");
EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 64), "64:ff9b::/64");
EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 127), "64:ff9b::16/127");
EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 128), "64:ff9b::17/128");
EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 0), "::/0");
EXPECT_THROW(getIPPrefixUsingVarchar("::ffff:1.2.3.4", -1), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("::ffff:1.2.3.4", 33), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("64:ff9b::10", -1), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("64:ff9b::10", 129), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("localhost", 24), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("64::ff9b::10", 24), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("64:face:book::10", 24), VeloxUserError);
EXPECT_THROW(getIPPrefixUsingVarchar("123.456.789.012", 24), VeloxUserError);
}


TEST_F(IPAddressTest, castRoundTripPrefix) {
auto strings = makeFlatVector<std::string>(
{"87a0:ce14:8989::/48", "7800::/5", "192.0.0.0/5"});
Expand Down Expand Up @@ -178,6 +267,15 @@ TEST_F(IPAddressTest, IPSubnetMin) {
getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"));
}

TEST_F(IPAddressTest, IPSubnetMinPrestoTests) {
EXPECT_EQ(getIPSubnetMin("1.2.3.4/24"), "1.2.3.0");
EXPECT_EQ(getIPSubnetMin("1.2.3.4/32"), "1.2.3.4");
EXPECT_EQ(getIPSubnetMin("64:ff9b::17/64"), "64:ff9b::");
EXPECT_EQ(getIPSubnetMin("64:ff9b::17/127"), "64:ff9b::16");
EXPECT_EQ(getIPSubnetMin("64:ff9b::17/128"), "64:ff9b::17");
EXPECT_EQ(getIPSubnetMin("64:ff9b::17/0"), "::");
}

TEST_F(IPAddressTest, IPSubnetMax) {
EXPECT_EQ("192.127.255.255", getIPSubnetMax("192.64.1.1/9"));
EXPECT_EQ("255.255.255.255", getIPSubnetMax("192.64.1.1/0"));
Expand All @@ -202,28 +300,39 @@ TEST_F(IPAddressTest, IPSubnetMax) {
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"));
}

TEST_F(IPAddressTest, IPSubnetMaxPrestoTests) {
EXPECT_EQ(getIPSubnetMax("1.2.3.128/26"),"1.2.3.191");
EXPECT_EQ(getIPSubnetMax("192.168.128.4/32"), "192.168.128.4");
EXPECT_EQ(getIPSubnetMax("10.1.16.3/9"), "10.127.255.255");
EXPECT_EQ(getIPSubnetMax("2001:db8::16/127"), "2001:db8::17");
EXPECT_EQ(getIPSubnetMax("2001:db8::16/128"), "2001:db8::16");
EXPECT_EQ(getIPSubnetMax("64:ff9b::17/64"), "64:ff9b::ffff:ffff:ffff:ffff");
EXPECT_EQ(getIPSubnetMax("64:ff9b::17/72"), "64:ff9b::ff:ffff:ffff:ffff");
EXPECT_EQ(getIPSubnetMax("64:ff9b::17/0"), "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff");
}

TEST_F(IPAddressTest, IPSubnetRange) {
EXPECT_EQ("192.0.0.0", getIPSubnetRangeMin("192.64.1.1/9"));
EXPECT_EQ("192.127.255.255", getIPSubnetMax("192.64.1.1/9"));
EXPECT_EQ("192.127.255.255", getIPSubnetRangeMax("192.64.1.1/9"));

EXPECT_EQ("0.0.0.0", getIPSubnetRangeMin("192.64.1.1/0"));
EXPECT_EQ("255.255.255.255", getIPSubnetMax("192.64.1.1/0"));
EXPECT_EQ("255.255.255.255", getIPSubnetRangeMax("192.64.1.1/0"));

EXPECT_EQ("128.0.0.0", getIPSubnetRangeMin("192.64.1.1/1"));
EXPECT_EQ("255.255.255.255", getIPSubnetMax("192.64.1.1/1"));
EXPECT_EQ("255.255.255.255", getIPSubnetRangeMax("192.64.1.1/1"));

EXPECT_EQ("192.64.1.0", getIPSubnetRangeMin("192.64.1.1/31"));
EXPECT_EQ("192.64.1.1", getIPSubnetMax("192.64.1.1/31"));
EXPECT_EQ("192.64.1.1", getIPSubnetRangeMax("192.64.1.1/31"));

EXPECT_EQ("192.64.1.1", getIPSubnetRangeMin("192.64.1.1/32"));
EXPECT_EQ("192.64.1.1", getIPSubnetMax("192.64.1.1/32"));
EXPECT_EQ("192.64.1.1", getIPSubnetRangeMax("192.64.1.1/32"));

EXPECT_EQ(
"2001:db8:85a3::",
getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48"));
EXPECT_EQ(
"2001:db8:85a3:ffff:ffff:ffff:ffff:ffff",
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48"));
getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48"));

EXPECT_EQ("::", getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0"));
EXPECT_EQ("::", getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1"));
Expand All @@ -236,16 +345,66 @@ TEST_F(IPAddressTest, IPSubnetRange) {

EXPECT_EQ(
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff",
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0"));
getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0"));
EXPECT_EQ(
"7fff:ffff:ffff:ffff:ffff:ffff:ffff:ffff",
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1"));
getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1"));
EXPECT_EQ(
"2001:db8:85a3:1:1:8a2e:370:7335",
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127"));
getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127"));
EXPECT_EQ(
"2001:db8:85a3:1:1:8a2e:370:7334",
getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"));
getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"));
}

TEST_F(IPAddressTest, IPSubnetRangePrestoTests){
EXPECT_EQ(getIPSubnetRangeMin("1.2.3.160/24"), "1.2.3.0");
EXPECT_EQ(getIPSubnetRangeMin("1.2.3.128/31"), "1.2.3.128");
EXPECT_EQ(getIPSubnetRangeMin("10.1.6.46/32"), "10.1.6.46");
EXPECT_EQ(getIPSubnetRangeMin("10.1.6.46/0"), "0.0.0.0");
EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::17/64"), "64:ff9b::");
EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::52f4/120"), "64:ff9b::5200");
EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::17/128"), "64:ff9b::17");

EXPECT_EQ(getIPSubnetRangeMax("1.2.3.160/24"), "1.2.3.255");
EXPECT_EQ(getIPSubnetRangeMax("1.2.3.128/31"), "1.2.3.129");
EXPECT_EQ(getIPSubnetRangeMax("10.1.6.46/32"), "10.1.6.46");
EXPECT_EQ(getIPSubnetRangeMax("10.1.6.46/0"), "255.255.255.255");
EXPECT_EQ(getIPSubnetRangeMax("64:ff9b::17/64"), "64:ff9b::ffff:ffff:ffff:ffff");
EXPECT_EQ(getIPSubnetRangeMax("64:ff9b::52f4/120"), "64:ff9b::52ff");
EXPECT_EQ(getIPSubnetRangeMax("64:ff9b::17/128"), "64:ff9b::17");
}

TEST_F(IPAddressTest, IPSubnetOfIPAddress) {
EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26","1.2.3.129"), true);
EXPECT_EQ(getIsSubnetOfIP("64:fa9b::17/64", "64:ffff::17"), false);
}

TEST_F(IPAddressTest, IPSubnetOfIPPrefix) {
EXPECT_EQ(getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false);
EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/32", "64:ffff::17/24"), false);
EXPECT_EQ(getIsSubnetOfIPPrefix("64:ffff::17/24", "64:ff9b::17/32"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.131/26"), true);
}

TEST_F(IPAddressTest, IPSubnetOfPrestoTests) {
EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26", "1.2.3.129"), true);
EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26", "1.2.5.1"), false);
EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/32", "1.2.3.128"), true);
EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/0", "192.168.5.1"), true);
EXPECT_EQ(getIsSubnetOfIP("64:ff9b::17/64", "64:ff9b::ffff:ff"), true);
EXPECT_EQ(getIsSubnetOfIP("64:ff9b::17/64", "64:ffff::17"), false);

EXPECT_EQ(getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("1.2.3.128/26", "1.2.5.1/30"), false);
EXPECT_EQ(getIsSubnetOfIPPrefix("1.2.3.128/26", "1.2.3.128/26"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ff9b::ff:25/80"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false);
EXPECT_EQ(getIsSubnetOfIPPrefix("2804:431:b000::/37", "2804:431:b000::/38"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("2804:431:b000::/38", "2804:431:b000::/37"), false);
EXPECT_EQ(getIsSubnetOfIPPrefix("170.0.52.0/22", "170.0.52.0/24"), true);
EXPECT_EQ(getIsSubnetOfIPPrefix("170.0.52.0/24", "170.0.52.0/22"), false);
}

} // namespace
Expand Down

0 comments on commit 21a8d16

Please sign in to comment.