From f37e0fd8121cd05afc3e0b704383fc103e777a26 Mon Sep 17 00:00:00 2001 From: Yenda Li Date: Wed, 12 Feb 2025 18:19:13 -0800 Subject: [PATCH] feat: Add ip_subnet_range and is_subnet_of [6/n] (#11777) Summary: Add function for ip_subnet_range and is_subnet_of. This is based on split of https://github.com/facebookincubator/velox/pull/11407 Differential Revision: D65839268 --- velox/docs/functions/presto/ipaddress.rst | 23 ++++ .../functions/prestosql/IPAddressFunctions.h | 95 ++++++++++++--- .../tests/IPAddressFunctionsTest.cpp | 112 ++++++++++++++++++ 3 files changed, 213 insertions(+), 17 deletions(-) diff --git a/velox/docs/functions/presto/ipaddress.rst b/velox/docs/functions/presto/ipaddress.rst index b497c1043dac..dcf5d4923b7b 100644 --- a/velox/docs/functions/presto/ipaddress.rst +++ b/velox/docs/functions/presto/ipaddress.rst @@ -26,3 +26,26 @@ IP Functions SELECT ip_subnet_max(IPPREFIX '192.64.0.0/9'); -- {192.127.255.255} SELECT ip_subnet_max(IPPREFIX '2001:0db8:85a3:0001:0001:8a2e:0370:7334/48'); -- {2001:db8:85a3:ffff:ffff:ffff:ffff:ffff} +.. function:: ip_subnet_range(ip_prefix) -> array(ip_address) + + Return an array of 2 IP addresses. + The array contains the smallest and the largest IP address + in the subnet specified by ``ip_prefix``. :: + + SELECT ip_subnet_range(IPPREFIX '1.2.3.160/24'); -- [{1.2.3.0}, {1.2.3.255}] + SELECT ip_subnet_range(IPPREFIX '64:ff9b::52f4/120'); -- [{64:ff9b::5200}, {64:ff9b::52ff}] + +.. function:: is_subnet_of(ip_prefix, ip_address) -> boolean + + Returns ``true`` if the ``ip_address`` is in the subnet of ``ip_prefix``. :: + + SELECT is_subnet_of(IPPREFIX '1.2.3.128/26', IPADDRESS '1.2.3.129'); -- true + SELECT is_subnet_of(IPPREFIX '64:fa9b::17/64', IPADDRESS '64:ffff::17'); -- false + +.. function:: is_subnet_of(ip_prefix1, ip_prefix2) -> boolean + + Returns ``true`` if ``ip_prefix2`` is a subnet of ``ip_prefix1``. :: + + SELECT is_subnet_of(IPPREFIX '192.168.3.131/26', IPPREFIX '192.168.3.144/30'); -- true + SELECT is_subnet_of(IPPREFIX '64:ff9b::17/64', IPPREFIX '64:ffff::17/64'); -- false + SELECT is_subnet_of(IPPREFIX '192.168.3.131/26', IPPREFIX '192.168.3.131/26'); -- true diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h index 2bd813c5ccf3..60002190e355 100644 --- a/velox/functions/prestosql/IPAddressFunctions.h +++ b/velox/functions/prestosql/IPAddressFunctions.h @@ -21,6 +21,32 @@ #include "velox/functions/prestosql/types/IPPrefixType.h" namespace facebook::velox::functions { +namespace { + +inline bool isIPv4(int128_t ip) { + int128_t ipV4 = 0x0000FFFF00000000; + uint128_t mask = 0xFFFFFFFFFFFFFFFF; + constexpr int kIPV6HalfBits = 64; + mask = (mask << kIPV6HalfBits) | 0xFFFFFFFF00000000; + return (ip & mask) == ipV4; +} + +inline int128_t getIPSubnetMax(int128_t ip, uint8_t prefix) { + uint128_t mask = 1; + if (isIPv4(ip)) { + ip |= (mask << (ipaddress::kIPV4Bits - prefix)) - 1; + return ip; + } + + // Special case: Overflow to all 0 subtracting 1 does not work. + if (prefix == 0) { + return -1; + } + + ip |= (mask << (ipaddress::kIPV6Bits - prefix)) - 1; + return ip; +} +} // namespace template struct IPPrefixFunction { @@ -97,28 +123,57 @@ struct IPSubnetMaxFunction { result = getIPSubnetMax(*ipPrefix.template at<0>(), *ipPrefix.template at<1>()); } +}; - private: - static int128_t getIPSubnetMax(int128_t ip, uint8_t prefix) { - auto tryIpv4 = ipaddress::tryIpPrefixLengthFromIPAddressType(ip); - // This check should never fail because we're taking a pre-existing - // IPPrefix. - VELOX_CHECK(tryIpv4.hasValue()); +template +struct IPSubnetRangeFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); - const bool isIpV4 = tryIpv4.value() == ipaddress::kIPV4Bits; - uint128_t mask = 1; - if (isIpV4) { - ip |= (mask << (ipaddress::kIPV4Bits - prefix)) - 1; - return ip; - } + FOLLY_ALWAYS_INLINE void call( + out_type>& result, + const arg_type& ipPrefix) { + result.push_back(*ipPrefix.template at<0>()); + result.push_back( + getIPSubnetMax(*ipPrefix.template at<0>(), *ipPrefix.template at<1>())); + } +}; + +template +struct IPSubnetOfFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix, + const arg_type& ip) { + result = isSubnetOf(ipPrefix, *ip); + } - // Special case: Overflow to all 0 subtracting 1 does not work. - if (prefix == 0) { - return -1; + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix, + const arg_type& ipPrefix2) { + result = (*ipPrefix2.template at<1>() >= *ipPrefix.template at<1>()) && + isSubnetOf(ipPrefix, *ipPrefix2.template at<0>()); + } + + private: + static bool isSubnetOf(const arg_type& ipPrefix, int128_t checkIP) { + uint128_t mask = 1; + const uint8_t prefix = *ipPrefix.template at<1>(); + if (isIPv4(*ipPrefix.template at<0>())) { + checkIP &= ((mask << (ipaddress::kIPV4Bits - prefix)) - 1) ^ + static_cast(-1); + } else { + // Special case: Overflow to all 0 subtracting 1 does not work. + if (prefix == 0) { + checkIP = 0; + } else { + checkIP &= ((mask << (ipaddress::kIPV6Bits - prefix)) - 1) ^ + static_cast(-1); + } } - ip |= (mask << (ipaddress::kIPV6Bits - prefix)) - 1; - return ip; + return (*ipPrefix.template at<0>() == checkIP); } }; @@ -133,6 +188,12 @@ void registerIPAddressFunctions(const std::string& prefix) { {prefix + "ip_subnet_min"}); registerFunction( {prefix + "ip_subnet_max"}); + registerFunction, IPPrefix>( + {prefix + "ip_subnet_range"}); + registerFunction( + {prefix + "is_subnet_of"}); + registerFunction( + {prefix + "is_subnet_of"}); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp index 95f1dd2f8010..383ff6467971 100644 --- a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp @@ -45,6 +45,36 @@ class IPAddressFunctionsTest : public functions::test::FunctionBaseTest { return evaluateOnce( "cast(ip_subnet_max(cast(c0 as ipprefix)) as varchar)", input); } + + std::optional ipSubnetRangeMin( + const std::optional& input) { + return evaluateOnce( + "cast(ip_subnet_range(cast(c0 as ipprefix))[1] as varchar)", input); + } + + std::optional ipSubnetRangeMax( + const std::optional& input) { + return evaluateOnce( + "cast(ip_subnet_range(cast(c0 as ipprefix))[2] as varchar)", input); + } + + std::optional isSubnetOfIP( + const std::optional& prefix, + const std::optional& ip) { + return evaluateOnce( + "is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipaddress))", + prefix, + ip); + } + + std::optional isSubnetOfIPPrefix( + const std::optional& prefix, + const std::optional& prefix2) { + return evaluateOnce( + "is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipprefix))", + prefix, + prefix2); + } }; TEST_F(IPAddressFunctionsTest, ipPrefixFromIpAddress) { @@ -180,4 +210,86 @@ TEST_F(IPAddressFunctionsTest, ipSubnetMax) { "2001:db8:85a3:1:1:8a2e:370:7334"); } +TEST_F(IPAddressFunctionsTest, IPSubnetRange) { + ASSERT_EQ("192.0.0.0", ipSubnetRangeMin("192.64.1.1/9")); + ASSERT_EQ("192.127.255.255", ipSubnetRangeMax("192.64.1.1/9")); + ASSERT_EQ("0.0.0.0", ipSubnetRangeMin("192.64.1.1/0")); + ASSERT_EQ("255.255.255.255", ipSubnetRangeMax("192.64.1.1/0")); + ASSERT_EQ("128.0.0.0", ipSubnetRangeMin("192.64.1.1/1")); + ASSERT_EQ("255.255.255.255", ipSubnetRangeMax("192.64.1.1/1")); + ASSERT_EQ("192.64.1.0", ipSubnetRangeMin("192.64.1.1/31")); + ASSERT_EQ("192.64.1.1", ipSubnetRangeMax("192.64.1.1/31")); + ASSERT_EQ("192.64.1.1", ipSubnetRangeMin("192.64.1.1/32")); + ASSERT_EQ("192.64.1.1", ipSubnetRangeMax("192.64.1.1/32")); + ASSERT_EQ( + "2001:db8:85a3::", + ipSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + ASSERT_EQ( + "2001:db8:85a3:ffff:ffff:ffff:ffff:ffff", + ipSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + ASSERT_EQ( + "::", ipSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + ASSERT_EQ( + "::", ipSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + ASSERT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + ipSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + ASSERT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + ipSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); + ASSERT_EQ( + "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + ipSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + ASSERT_EQ( + "7fff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + ipSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + ASSERT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7335", + ipSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + ASSERT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + ipSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); + ASSERT_EQ("1.2.3.0", ipSubnetRangeMin("1.2.3.160/24")); + ASSERT_EQ("1.2.3.128", ipSubnetRangeMin("1.2.3.128/31")); + ASSERT_EQ("10.1.6.46", ipSubnetRangeMin("10.1.6.46/32")); + ASSERT_EQ("0.0.0.0", ipSubnetRangeMin("10.1.6.46/0")); + ASSERT_EQ("64:ff9b::", ipSubnetRangeMin("64:ff9b::17/64")); + ASSERT_EQ("64:ff9b::5200", ipSubnetRangeMin("64:ff9b::52f4/120")); + ASSERT_EQ("64:ff9b::17", ipSubnetRangeMin("64:ff9b::17/128")); + ASSERT_EQ("1.2.3.255", ipSubnetRangeMax("1.2.3.160/24")); + ASSERT_EQ("1.2.3.129", ipSubnetRangeMax("1.2.3.128/31")); + ASSERT_EQ("10.1.6.46", ipSubnetRangeMax("10.1.6.46/32")); + ASSERT_EQ("255.255.255.255", ipSubnetRangeMax("10.1.6.46/0")); + ASSERT_EQ("64:ff9b::ffff:ffff:ffff:ffff", ipSubnetRangeMax("64:ff9b::17/64")); + ASSERT_EQ("64:ff9b::52ff", ipSubnetRangeMax("64:ff9b::52f4/120")); + ASSERT_EQ("64:ff9b::17", ipSubnetRangeMax("64:ff9b::17/128")); +} + +TEST_F(IPAddressFunctionsTest, IPSubnetOfIPPrefix) { + EXPECT_EQ(isSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true); + EXPECT_EQ(isSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false); + EXPECT_EQ(isSubnetOfIPPrefix("64:ff9b::17/32", "64:ffff::17/24"), false); + EXPECT_EQ(isSubnetOfIPPrefix("64:ffff::17/24", "64:ff9b::17/32"), true); + EXPECT_EQ(isSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.131/26"), true); + + EXPECT_EQ(isSubnetOfIP("1.2.3.128/26", "1.2.3.129"), true); + EXPECT_EQ(isSubnetOfIP("1.2.3.128/26", "1.2.5.1"), false); + EXPECT_EQ(isSubnetOfIP("1.2.3.128/32", "1.2.3.128"), true); + EXPECT_EQ(isSubnetOfIP("1.2.3.128/0", "192.168.5.1"), true); + EXPECT_EQ(isSubnetOfIP("64:ff9b::17/64", "64:ff9b::ffff:ff"), true); + EXPECT_EQ(isSubnetOfIP("64:ff9b::17/64", "64:ffff::17"), false); + + EXPECT_EQ(isSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true); + EXPECT_EQ(isSubnetOfIPPrefix("1.2.3.128/26", "1.2.5.1/30"), false); + EXPECT_EQ(isSubnetOfIPPrefix("1.2.3.128/26", "1.2.3.128/26"), true); + EXPECT_EQ(isSubnetOfIPPrefix("64:ff9b::17/64", "64:ff9b::ff:25/80"), true); + EXPECT_EQ(isSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false); + EXPECT_EQ( + isSubnetOfIPPrefix("2804:431:b000::/37", "2804:431:b000::/38"), true); + EXPECT_EQ( + isSubnetOfIPPrefix("2804:431:b000::/38", "2804:431:b000::/37"), false); + EXPECT_EQ(isSubnetOfIPPrefix("170.0.52.0/22", "170.0.52.0/24"), true); + EXPECT_EQ(isSubnetOfIPPrefix("170.0.52.0/24", "170.0.52.0/22"), false); +} + } // namespace facebook::velox::functions::prestosql