22
33#include <ydb/library/yql/public/udf/udf_helpers.h>
44
5- #include < util/draft/ip.h>
5+ #include <library/cpp/ipv6_address/ipv6_address.h>
6+ #include <library/cpp/ipmath/ipmath.h>
67#include <util/generic/buffer.h>
78
89namespace {
@@ -13,39 +14,77 @@ namespace {
1314 using TUnboxedValue = NKikimr ::NUdf ::TUnboxedValue ;
1415 using TUnboxedValuePod = NKikimr ::NUdf ::TUnboxedValuePod ;
1516
16- struct TSerializeIpVisitor {
17- TStringRef operator ()(const TIp4& ip) const {
18- return TStringRef (reinterpret_cast <const char *>(&ip), 4 );
17+ struct TRawIp4 {
18+ ui8 a , b , c , d ;
19+ };
20+
21+ struct TRawIp6 {
22+ ui8 a1 , a0 , b1 , b0 , c1 , c0 , d1 , d0 , e1 , e0 , f1 , f0 , g1 , g0 , h1 , h0 ;
23+ };
24+
25+ TIpv6Address DeserializeAddress (const TStringRef & str ) {
26+ TIpv6Address addr ;
27+ if (str .Size () == 4 ) {
28+ TRawIp4 addr4 ;
29+ memcpy (& addr4 , str .Data (), sizeof addr4 );
30+ addr = {addr4 .a , addr4 .b , addr4 .c , addr4 .d };
31+ } else if (str .Size () == 16 ) {
32+ TRawIp6 addr6 ;
33+ memcpy (& addr6 , str .Data (), sizeof addr6 );
34+ addr = {ui16 (ui32 (addr6 .a1 ) << ui32 (8 ) | ui32 (addr6 .a0 )),
35+ ui16 (ui32 (addr6 .b1 ) << ui32 (8 ) | ui32 (addr6 .b0 )),
36+ ui16 (ui32 (addr6 .c1 ) << ui32 (8 ) | ui32 (addr6 .c0 )),
37+ ui16 (ui32 (addr6 .d1 ) << ui32 (8 ) | ui32 (addr6 .d0 )),
38+ ui16 (ui32 (addr6 .e1 ) << ui32 (8 ) | ui32 (addr6 .e0 )),
39+ ui16 (ui32 (addr6 .f1 ) << ui32 (8 ) | ui32 (addr6 .f0 )),
40+ ui16 (ui32 (addr6 .g1 ) << ui32 (8 ) | ui32 (addr6 .g0 )),
41+ ui16 (ui32 (addr6 .h1 ) << ui32 (8 ) | ui32 (addr6 .h0 )),
42+ };
43+ } else {
44+ ythrow yexception () << "Incorrect size of input, expected "
45+ << "4 or 16, got " << str .Size ();
1946 }
20- TStringRef operator ()(const TIp6& ip) const {
21- return TStringRef (reinterpret_cast <const char *>(&ip.Data ), 16 );
47+ return addr ;
48+ }
49+
50+ TString SerializeAddress (const TIpv6Address & addr ) {
51+ Y_ENSURE (addr .Type () == TIpv6Address ::Ipv4 || addr .Type () == TIpv6Address ::Ipv6 );
52+ TString res ;
53+ ui128 x = addr ;
54+ if (addr .Type () == TIpv6Address ::Ipv4 ) {
55+ TRawIp4 addr4 {
56+ ui8 (x >> 24 & 0xff ),
57+ ui8 (x >> 16 & 0xff ),
58+ ui8 (x >> 8 & 0xff ),
59+ ui8 (x & 0xff )
60+ };
61+ res = TString (reinterpret_cast < const char * > (& addr4 ), sizeof addr4 );
62+ } else if (addr .Type () == TIpv6Address ::Ipv6 ) {
63+ TRawIp6 addr6 {
64+ ui8 (x >> 120 & 0xff ), ui8 (x >> 112 & 0xff ),
65+ ui8 (x >> 104 & 0xff ), ui8 (x >> 96 & 0xff ),
66+ ui8 (x >> 88 & 0xff ), ui8 (x >> 80 & 0xff ),
67+ ui8 (x >> 72 & 0xff ), ui8 (x >> 64 & 0xff ),
68+ ui8 (x >> 56 & 0xff ), ui8 (x >> 48 & 0xff ),
69+ ui8 (x >> 40 & 0xff ), ui8 (x >> 32 & 0xff ),
70+ ui8 (x >> 24 & 0xff ), ui8 (x >> 16 & 0xff ),
71+ ui8 (x >> 8 & 0xff ), ui8 (x & 0xff )
72+ };
73+ res = TString (reinterpret_cast < const char * > (& addr6 ), sizeof addr6 );
2274 }
23- };
75+ return res ;
76+ }
2477
2578 SIMPLE_STRICT_UDF (TFromString , TOptionalString (TAutoMapString )) {
26- try {
27- TString input (args[0 ].AsStringRef ());
28- const TIp4Or6& ip = Ip4Or6FromString (input.c_str ());
29- return valueBuilder->NewString (std::visit (TSerializeIpVisitor (), ip));
30- } catch (TSystemError&) {
79+ TIpv6Address addr = TIpv6Address ::FromString (args [0 ].AsStringRef ());
80+ if (addr .Type () != TIpv6Address ::Ipv4 && addr .Type () != TIpv6Address ::Ipv6 ) {
3181 return TUnboxedValue ();
3282 }
83+ return valueBuilder -> NewString (SerializeAddress (addr ));
3384 }
3485
3586 SIMPLE_UDF (TToString , char * (TAutoMapString )) {
36- const auto & ref = args[0 ].AsStringRef ();
37- if (ref.Size () == 4 ) {
38- TIp4 ip;
39- memcpy (&ip, ref.Data (), sizeof (ip));
40- return valueBuilder->NewString (Ip4Or6ToString (ip));
41- } else if (ref.Size () == 16 ) {
42- TIp6 ip;
43- memcpy (&ip.Data , ref.Data (), sizeof (ip.Data ));
44- return valueBuilder->NewString (Ip4Or6ToString (ip));
45- } else {
46- ythrow yexception () << " Incorrect size of input, expected "
47- << " 4 or 16, got " << ref.Size ();
48- }
87+ return valueBuilder -> NewString (DeserializeAddress (args [0 ].AsStringRef ()).ToString (false));
4988 }
5089
5190 SIMPLE_STRICT_UDF (TIsIPv4 , bool (TOptionalString )) {
@@ -73,15 +112,8 @@ namespace {
73112 bool result = false;
74113 if (args [0 ]) {
75114 const auto ref = args [0 ].AsStringRef ();
76- if (ref.Size () == 16 && ref.Data ()[10 ] == -1 ) {
77- bool allZeroes = true ;
78- for (int i = 0 ; i < 10 ; ++i) {
79- if (ref.Data ()[i] != 0 ) {
80- allZeroes = false ;
81- break ;
82- }
83- }
84- result = allZeroes;
115+ if (ref .Size () == 16 ) {
116+ result = DeserializeAddress (ref ).Isv4MappedTov6 ();
85117 }
86118 }
87119 return TUnboxedValuePod (result );
@@ -92,10 +124,9 @@ namespace {
92124 if (ref .Size () == 16 ) {
93125 return valueBuilder -> NewString (ref );
94126 } else if (ref .Size () == 4 ) {
95- TIp4 ipv4;
96- memcpy (&ipv4, ref.Data (), sizeof (ipv4));
97- const TIp6 ipv6 = Ip6FromIp4 (ipv4);
98- return valueBuilder->NewString (TStringRef (reinterpret_cast <const char *>(&ipv6.Data ), 16 ));
127+ TIpv6Address addr4 = DeserializeAddress (ref );
128+ auto addr6 = TIpv6Address (ui128 (addr4 ) | ui128 (0xFFFF ) << 32 , TIpv6Address ::Ipv6 );
129+ return valueBuilder -> NewString (SerializeAddress (addr6 ));
99130 } else {
100131 ythrow yexception () << "Incorrect size of input, expected "
101132 << "4 or 16, got " << ref .Size ();
@@ -105,33 +136,27 @@ namespace {
105136 SIMPLE_UDF_WITH_OPTIONAL_ARGS (TGetSubnet , char * (TAutoMapString , TOptionalByte ), 1 ) {
106137 const auto ref = args [0 ].AsStringRef ();
107138 ui8 subnetSize = args [1 ].GetOrDefault < ui8 > (0 );
108-
139+ TIpv6Address addr = DeserializeAddress ( ref );
109140 if (ref .Size () == 4 ) {
110141 if (!subnetSize ) {
111142 subnetSize = 24 ;
112143 }
144+ if (subnetSize > 32 ) {
145+ subnetSize = 32 ;
146+ }
113147 } else if (ref .Size () == 16 ) {
114148 if (!subnetSize ) {
115149 subnetSize = 64 ;
116150 }
151+ if (subnetSize > 128 ) {
152+ subnetSize = 128 ;
153+ }
117154 } else {
118155 ythrow yexception () << "Incorrect size of input, expected "
119156 << "4 or 16, got " << ref .Size ();
120157 }
121- TBuffer result (ref.Data (), ref.Size ());
122- int bytesToMask = ref.Size () * 8 - subnetSize;
123- ui8 currentByte = ref.Size () - 1 ;
124- while (bytesToMask > 0 ) {
125- if (bytesToMask > 8 ) {
126- result.Data ()[currentByte] = 0 ;
127- } else {
128- result.Data ()[currentByte] = result.Data ()[currentByte] & (0xff << bytesToMask);
129- }
130- bytesToMask -= 8 ;
131- --currentByte;
132- }
133-
134- return valueBuilder->NewString (TStringRef (result.Data (), result.Size ()));
158+ TIpv6Address beg = LowerBoundForPrefix (addr , subnetSize );
159+ return valueBuilder -> NewString (SerializeAddress (beg ));
135160 }
136161
137162#define EXPORTED_IP_BASE_UDF \
0 commit comments