Skip to content

Commit 21bfb72

Browse files
author
Vadim Averin
authored
YDB-2321 Refactor IP UDF to use supported ipmath library (#618)
* Apply patch with changes * Add & canonize some tests * More refactoring * Final cleanup * Apply patch with changes * Add & canonize some tests * More refactoring * Final cleanup * Add explicit dependencies on library/cpp/ipv6_address
1 parent b07e92f commit 21bfb72

File tree

5 files changed

+177
-54
lines changed

5 files changed

+177
-54
lines changed

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

+78-53
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
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

89
namespace {
@@ -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 \

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

+2
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ SRCS(
1212

1313
PEERDIR(
1414
ydb/library/yql/public/udf
15+
library/cpp/ipmath
16+
library/cpp/ipv6_address
1517
)
1618

1719
END()

‎ydb/library/yql/udfs/common/ip_base/test/canondata/test.test_Basic_/results.txt

+93
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,26 @@
8787
"String"
8888
]
8989
]
90+
];
91+
[
92+
"single_subnet4";
93+
[
94+
"OptionalType";
95+
[
96+
"DataType";
97+
"String"
98+
]
99+
]
100+
];
101+
[
102+
"single_subnet6";
103+
[
104+
"OptionalType";
105+
[
106+
"DataType";
107+
"String"
108+
]
109+
]
90110
]
91111
]
92112
]
@@ -113,6 +133,12 @@
113133
];
114134
[
115135
"127.0.0.0"
136+
];
137+
[
138+
"127.0.0.1"
139+
];
140+
[
141+
"127.0.0.1"
116142
]
117143
];
118144
[
@@ -136,6 +162,12 @@
136162
];
137163
[
138164
"::"
165+
];
166+
[
167+
"::"
168+
];
169+
[
170+
"::1"
139171
]
140172
];
141173
[
@@ -161,6 +193,12 @@
161193
];
162194
[
163195
"213.180.0.0"
196+
];
197+
[
198+
"213.180.193.3"
199+
];
200+
[
201+
"213.180.193.3"
164202
]
165203
];
166204
[
@@ -186,6 +224,12 @@
186224
];
187225
[
188226
"2a02::"
227+
];
228+
[
229+
"2a02:6b8::"
230+
];
231+
[
232+
"2a02:6b8::3"
189233
]
190234
];
191235
[
@@ -211,6 +255,12 @@
211255
];
212256
[
213257
"2400::"
258+
];
259+
[
260+
"2400:cb00::"
261+
];
262+
[
263+
"2400:cb00:2048:1::681c:1b65"
214264
]
215265
];
216266
[
@@ -236,6 +286,12 @@
236286
];
237287
[
238288
"fe80::"
289+
];
290+
[
291+
"fe80::"
292+
];
293+
[
294+
"fe80::215:b2ff:fea9:67ce"
239295
]
240296
];
241297
[
@@ -261,6 +317,12 @@
261317
];
262318
[
263319
"::"
320+
];
321+
[
322+
"::"
323+
];
324+
[
325+
"::ffff:77.75.155.3"
264326
]
265327
];
266328
[
@@ -272,7 +334,38 @@
272334
#;
273335
#;
274336
#;
337+
#;
338+
#;
275339
#
340+
];
341+
[
342+
[
343+
"\0\0\0\0"
344+
];
345+
[
346+
"0.0.0.0"
347+
];
348+
%true;
349+
%false;
350+
%false;
351+
[
352+
"::ffff:0.0.0.0"
353+
];
354+
[
355+
"0.0.0.0"
356+
];
357+
[
358+
"0.0.0.0"
359+
];
360+
[
361+
"0.0.0.0"
362+
];
363+
[
364+
"0.0.0.0"
365+
];
366+
[
367+
"0.0.0.0"
368+
]
276369
]
277370
]
278371
}

‎ydb/library/yql/udfs/common/ip_base/test/cases/Basic.in

+1
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,4 @@
66
{"key"="fe80::215:b2ff:fea9:67ce";"subkey"="";"value"=""};
77
{"key"="::ffff:77.75.155.3";"subkey"="";"value"=""};
88
{"key"="sdfsdfsdf";"subkey"="";"value"=""};
9+
{"key"="0.0.0.0";"subkey"="";value=""};

‎ydb/library/yql/udfs/common/ip_base/test/cases/Basic.sql

+3-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,9 @@ SELECT
88
Ip::ToString(Ip::ConvertToIPv6(internal_representation)) AS all_ipv6,
99
Ip::ToString(Ip::GetSubnet(internal_representation)) AS default_subnet,
1010
Ip::ToString(Ip::GetSubnet(internal_representation, 125)) AS small_subnet,
11-
Ip::ToString(Ip::GetSubnet(internal_representation, 16)) AS large_subnet
11+
Ip::ToString(Ip::GetSubnet(internal_representation, 16)) AS large_subnet,
12+
Ip::ToString(Ip::GetSubnet(internal_representation, 32)) AS single_subnet4,
13+
Ip::ToString(Ip::GetSubnet(internal_representation, 128)) AS single_subnet6
1214
FROM (
1315
SELECT Ip::FromString(key) AS internal_representation FROM Input
1416
);

0 commit comments

Comments
 (0)