Fix issue with ipv4->ipv6

mohsaka · Jun 28, 2024 · 6e11299 · 6e11299
1 parent 7886efe
commit 6e11299
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 11 deletions.
diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp
@@ -23,6 +23,21 @@ namespace {
 
 class IPAddressTest : public functions::test::FunctionBaseTest {
  protected:
+
+   std::optional<std::string> castIPAddress(
+      const std::optional<std::string> input) {
+    auto result = evaluateOnce<std::string>(
+        "cast(cast(c0 as ipaddress) as varchar)", input);
+    return result;
+  }
+
+   std::optional<std::string> castIPPrefix(
+      const std::optional<std::string> input) {
+    auto result = evaluateOnce<std::string>(
+        "cast(cast(c0 as ipprefix) as varchar)", input);
+    return result;
+  }
+
   std::optional<std::string> getIPPrefix(
       const std::optional<std::string> input,
       std::optional<int8_t> mask) {
@@ -88,17 +103,31 @@ class IPAddressTest : public functions::test::FunctionBaseTest {
   }
 };
 
-TEST_F(IPAddressTest, castAsVarchar) {
-  auto result = evaluate<FlatVector<StringView>>(
-      "cast(ipaddress() as varchar)", makeRowVector(ROW({}), 10));
+TEST_F(IPAddressTest, castFail) {
+  EXPECT_THROW(castIPAddress("12.483.09.1"), VeloxUserError);
+  EXPECT_THROW(castIPAddress("10.135.23.12.12"), VeloxUserError);
+  EXPECT_THROW(castIPAddress("10.135.23"), VeloxUserError);
+  EXPECT_THROW(castIPAddress("q001:0db8:85a3:0001:0001:8a2e:0370:7334"), VeloxUserError);
+  EXPECT_THROW(castIPAddress("2001:0db8:85a3:542e:0001:0001:8a2e:0370:7334"), VeloxUserError);
+  EXPECT_THROW(castIPAddress("2001:0db8:85a3:0001:0001:8a2e:0370"), VeloxUserError);
+
+  EXPECT_THROW(castIPPrefix("12.135.23.12/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("10.135.23.12/33"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/33"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64:ff9b::10/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64:ff9b::10/129"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/33"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:909:909/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("::ffff:909:909/33"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64:ff9b::10/-1"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64:ff9b::10/129"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("localhost/24"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64::ff9b::10/24"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("64:face:book::10/24"), VeloxUserError);
+  EXPECT_THROW(castIPPrefix("123.456.789.012/24"), VeloxUserError);
 
-  // Sanity check results. All strings are unique.
-  std::unordered_set<std::string> ipaddresses;
-  for (auto i = 0; i < 10; ++i) {
-    const auto ipaddress = result->valueAt(i).str();
-    ASSERT_TRUE(ipaddresses.insert(ipaddress).second);
-  }
-  ASSERT_EQ(10, ipaddresses.size());
 }
 
 TEST_F(IPAddressTest, castRoundTrip) {

diff --git a/velox/functions/prestosql/types/IPPrefixType.cpp b/velox/functions/prestosql/types/IPPrefixType.cpp
@@ -125,8 +125,16 @@ class IPPrefixCastOperator : public exec::CastOperator {
                         .to_bytes();
       } else {
         auto v6Net = boost::asio::ip::make_network_v6(ipAddressString);
+        if(addr.to_v6().is_v4_mapped()){
+          auto v4Addr = boost::asio::ip::make_address_v4(boost::asio::ip::v4_mapped, addr.to_v6());
+          auto v4Net = boost::asio::ip::make_network_v4(v4Addr, (uint8_t)v6Net.prefix_length());
+          addrBytes = boost::asio::ip::make_address_v6(
+                        boost::asio::ip::v4_mapped, v4Net.canonical().address())
+                        .to_bytes();
+        }else{
+          addrBytes = v6Net.canonical().address().to_bytes();
+        }
         res.prefix = (uint8_t)v6Net.prefix_length();
-        addrBytes = v6Net.canonical().address().to_bytes();
       }
 
       bigEndianByteArray(addrBytes);