Apply patch with changes

Author: Vadim Averin

ydb/library/yql/udfs/common/ip_base/lib/ip_base_udf.h

@@ -2,7 +2,7 @@
 
 #include <ydb/library/yql/public/udf/udf_helpers.h>
 
-#include <util/draft/ip.h>
+#include <library/cpp/ipmath/ipmath.h>
 #include <util/generic/buffer.h>
 
 namespace {
@@ -13,39 +13,76 @@ namespace {
     using TUnboxedValue = NKikimr::NUdf::TUnboxedValue;
     using TUnboxedValuePod = NKikimr::NUdf::TUnboxedValuePod;
 
-    struct TSerializeIpVisitor {
-        TStringRef operator()(const TIp4& ip) const {
-            return TStringRef(reinterpret_cast<const char*>(&ip), 4);
-        }
-        TStringRef operator()(const TIp6& ip) const {
-            return TStringRef(reinterpret_cast<const char*>(&ip.Data), 16);
-        }
+    struct TRawIp4 {
+        ui8 a, b, c, d;
     };
 
-    SIMPLE_STRICT_UDF(TFromString, TOptionalString(TAutoMapString)) {
-        try {
-            TString input(args[0].AsStringRef());
-            const TIp4Or6& ip = Ip4Or6FromString(input.c_str());
-            return valueBuilder->NewString(std::visit(TSerializeIpVisitor(), ip));
-        } catch (TSystemError&) {
-            return TUnboxedValue();
+    struct TRawIp6 {
+        ui8 a1, a0, b1, b0, c1, c0, d1, d0, e1, e0, f1, f0, g1, g0, h1, h0;
+    };
+
+    TIpv6Address DeserializeAddress(const TStringRef& str) {
+        TIpv6Address addr;
+        if (str.Size() == 4) {
+            TRawIp4 addr4;
+            memcpy(&addr4, str.Data(), sizeof addr4);
+            addr = {addr4.a, addr4.b, addr4.c, addr4.d};
+        } else if (str.Size() == 16) {
+            TRawIp6 addr6;
+            memcpy(&addr6, str.Data(), sizeof addr6);
+            addr = {ui16(ui32(addr6.a1) << ui32(8) | ui32(addr6.a0)),
+                    ui16(ui32(addr6.b1) << ui32(8) | ui32(addr6.b0)),
+                    ui16(ui32(addr6.c1) << ui32(8) | ui32(addr6.c0)),
+                    ui16(ui32(addr6.d1) << ui32(8) | ui32(addr6.d0)),
+                    ui16(ui32(addr6.e1) << ui32(8) | ui32(addr6.e0)),
+                    ui16(ui32(addr6.f1) << ui32(8) | ui32(addr6.f0)),
+                    ui16(ui32(addr6.g1) << ui32(8) | ui32(addr6.g0)),
+                    ui16(ui32(addr6.h1) << ui32(8) | ui32(addr6.h0)),
+            };
+        } else {
+            ythrow yexception() << "Incorrect size of input, expected "
+            << "4 or 16, got " << str.Size();
         }
+        return addr;
     }
 
-    SIMPLE_UDF(TToString, char*(TAutoMapString)) {
-        const auto& ref = args[0].AsStringRef();
-        if (ref.Size() == 4) {
-            TIp4 ip;
-            memcpy(&ip, ref.Data(), sizeof(ip));
-            return valueBuilder->NewString(Ip4Or6ToString(ip));
-        } else if (ref.Size() == 16) {
-            TIp6 ip;
-            memcpy(&ip.Data, ref.Data(), sizeof(ip.Data));
-            return valueBuilder->NewString(Ip4Or6ToString(ip));
+    TMaybe<TString> SerializeAddress(const TIpv6Address& addr) {
+        TString res;
+        ui128 x = addr;
+        if (addr.Type() == TIpv6Address::Ipv4) {
+            TRawIp4 addr4 {
+                ui8(x >> 24 & 0xff),
+                ui8(x >> 16 & 0xff),
+                ui8(x >> 8  & 0xff),
+                ui8(x & 0xff)
+            };
+            res = TString(reinterpret_cast<const char *>(&addr4), sizeof addr4);
+        } else if (addr.Type() == TIpv6Address::Ipv6) {
+            TRawIp6 addr6 {
+                ui8(x >> 120 & 0xff), ui8(x >> 112 & 0xff),
+                ui8(x >> 104 & 0xff), ui8(x >> 96 & 0xff),
+                ui8(x >> 88 & 0xff), ui8(x >> 80 & 0xff),
+                ui8(x >> 72 & 0xff), ui8(x >> 64 & 0xff),
+                ui8(x >> 56 & 0xff), ui8(x >> 48 & 0xff),
+                ui8(x >> 40 & 0xff), ui8(x >> 32 & 0xff),
+                ui8(x >> 24 & 0xff), ui8(x >> 16 & 0xff),
+                ui8(x >> 8 & 0xff), ui8(x & 0xff)
+            };
+            res = TString(reinterpret_cast<const char *>(&addr6), sizeof addr6);
         } else {
-            ythrow yexception() << "Incorrect size of input, expected "
-            << "4 or 16, got " << ref.Size();
+            return Nothing();
         }
+        return res;
+    }
+
+    SIMPLE_STRICT_UDF(TFromString, TOptionalString(TAutoMapString)) {
+        TIpv6Address addr = TIpv6Address::FromString(args[0].AsStringRef());
+        auto res = SerializeAddress(addr);
+        return res ? valueBuilder->NewString(res.GetRef()) : TUnboxedValue(TUnboxedValuePod());
+    }
+
+    SIMPLE_UDF(TToString, char*(TAutoMapString)) {
+        return valueBuilder->NewString(DeserializeAddress(args[0].AsStringRef()).ToString(false));
     }
 
     SIMPLE_STRICT_UDF(TIsIPv4, bool(TOptionalString)) {
@@ -73,15 +110,8 @@ namespace {
         bool result = false;
         if (args[0]) {
             const auto ref = args[0].AsStringRef();
-            if (ref.Size() == 16 && ref.Data()[10] == -1) {
-                bool allZeroes = true;
-                for (int i = 0; i < 10; ++i) {
-                    if (ref.Data()[i] != 0) {
-                        allZeroes = false;
-                        break;
-                    }
-                }
-                result = allZeroes;
+            if (ref.Size() == 16) {
+                result = DeserializeAddress(ref).Isv4MappedTov6();
             }
         }
         return TUnboxedValuePod(result);
@@ -92,10 +122,8 @@ namespace {
         if (ref.Size() == 16) {
             return valueBuilder->NewString(ref);
         } else if (ref.Size() == 4) {
-            TIp4 ipv4;
-            memcpy(&ipv4, ref.Data(), sizeof(ipv4));
-            const TIp6 ipv6 = Ip6FromIp4(ipv4);
-            return valueBuilder->NewString(TStringRef(reinterpret_cast<const char*>(&ipv6.Data), 16));
+            auto addr6 = TIpv6Address::FromString("::ffff:" + DeserializeAddress(ref).ToString(false));
+            return valueBuilder->NewString(SerializeAddress(addr6).GetRef());
         } else {
             ythrow yexception() << "Incorrect size of input, expected "
             << "4 or 16, got " << ref.Size();
@@ -105,33 +133,28 @@ namespace {
     SIMPLE_UDF_WITH_OPTIONAL_ARGS(TGetSubnet, char*(TAutoMapString, TOptionalByte), 1) {
         const auto ref = args[0].AsStringRef();
         ui8 subnetSize = args[1].GetOrDefault<ui8>(0);
-
+        TIpv6Address addr = DeserializeAddress(ref);
         if (ref.Size() == 4) {
             if (!subnetSize) {
                 subnetSize = 24;
             }
+            if (subnetSize > 32) {
+                subnetSize = 32;
+            }
         } else if (ref.Size() == 16) {
             if (!subnetSize) {
                 subnetSize = 64;
             }
+            if (subnetSize > 128) {
+                subnetSize = 128;
+            }
         } else {
             ythrow yexception() << "Incorrect size of input, expected "
             << "4 or 16, got " << ref.Size();
         }
-        TBuffer result(ref.Data(), ref.Size());
-        int bytesToMask = ref.Size() * 8 - subnetSize;
-        ui8 currentByte = ref.Size() - 1;
-        while (bytesToMask > 0) {
-            if (bytesToMask > 8) {
-                result.Data()[currentByte] = 0;
-            } else {
-                result.Data()[currentByte] = result.Data()[currentByte] & (0xff << bytesToMask);
-            }
-            bytesToMask -= 8;
-            --currentByte;
-        }
-
-        return valueBuilder->NewString(TStringRef(result.Data(), result.Size()));
+        ui128 mask = ui128(-1) << int((addr.Type() == TIpv6Address::Ipv6 ? 128 : 32) - subnetSize);
+        TIpv6Address beg(ui128(addr) & mask, addr.Type());
+        return valueBuilder->NewString(SerializeAddress(beg).GetRef());
     }
 
 #define EXPORTED_IP_BASE_UDF \

ydb/library/yql/udfs/common/ip_base/lib/ya.make

@@ -12,6 +12,7 @@ SRCS(
 
 PEERDIR(
     ydb/library/yql/public/udf
+    library/cpp/ipmath
 )
 
 END()