ClickHouse

1
#include <Columns/ColumnConst.h>
2
#include <Columns/ColumnString.h>
3
#include <Columns/ColumnsNumber.h>
4
#include <Common/IPv6ToBinary.h>
5
#include <Common/formatIPv6.h>
6
#include <DataTypes/DataTypeNullable.h>
7
#include <DataTypes/DataTypesNumber.h>
8
#include <Functions/IFunction.h>
9
#include <Functions/FunctionFactory.h>
10
#include <Functions/FunctionHelpers.h>
11
#include <variant>
12
#include <charconv>
13

14

15
#include <Common/logger_useful.h>
16
namespace DB::ErrorCodes
17
{
18
    extern const int CANNOT_PARSE_TEXT;
19
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
20
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
21
}
22

23
namespace
24
{
25

26
class IPAddressVariant
27
{
28
public:
29

30
    explicit IPAddressVariant(std::string_view address_str)
31
    {
32
        UInt32 v4;
33
        if (DB::parseIPv4whole(address_str.begin(), address_str.end(), reinterpret_cast<unsigned char *>(&v4)))
34
        {
35
            addr = v4;
36
        }
37
        else
38
        {
39
            addr = IPv6AddrType();
40
            bool success = DB::parseIPv6whole(address_str.begin(), address_str.end(), std::get<IPv6AddrType>(addr).data());
41
            if (!success)
42
                throw DB::Exception(DB::ErrorCodes::CANNOT_PARSE_TEXT, "Neither IPv4 nor IPv6 address: '{}'", address_str);
43
        }
44
    }
45

46
    UInt32 asV4() const
47
    {
48
        if (const auto * val = std::get_if<IPv4AddrType>(&addr))
49
            return *val;
50
        return 0;
51
    }
52

53
    const uint8_t * asV6() const
54
    {
55
        if (const auto * val = std::get_if<IPv6AddrType>(&addr))
56
            return val->data();
57
        return nullptr;
58
    }
59

60
private:
61
    using IPv4AddrType = UInt32;
62
    using IPv6AddrType = std::array<uint8_t, IPV6_BINARY_LENGTH>;
63

64
    std::variant<IPv4AddrType, IPv6AddrType> addr;
65
};
66

67
struct IPAddressCIDR
68
{
69
    IPAddressVariant address;
70
    UInt8 prefix;
71
};
72

73
IPAddressCIDR parseIPWithCIDR(std::string_view cidr_str)
74
{
75
    size_t pos_slash = cidr_str.find('/');
76

77
    if (pos_slash == 0)
78
        throw DB::Exception(DB::ErrorCodes::CANNOT_PARSE_TEXT, "Error parsing IP address with prefix: {}", std::string(cidr_str));
79
    if (pos_slash == std::string_view::npos)
80
        throw DB::Exception(DB::ErrorCodes::CANNOT_PARSE_TEXT, "The text does not contain '/': {}", std::string(cidr_str));
81

82
    std::string_view addr_str = cidr_str.substr(0, pos_slash);
83
    IPAddressVariant addr(addr_str);
84

85
    uint8_t prefix = 0;
86
    auto prefix_str = cidr_str.substr(pos_slash+1);
87

88
    const auto * prefix_str_end = prefix_str.data() + prefix_str.size();
89
    auto [parse_end, parse_error] = std::from_chars(prefix_str.data(), prefix_str_end, prefix);
90
    uint8_t max_prefix = (addr.asV6() ? IPV6_BINARY_LENGTH : IPV4_BINARY_LENGTH) * 8;
91
    bool has_error = parse_error != std::errc() || parse_end != prefix_str_end || prefix > max_prefix;
92
    if (has_error)
93
        throw DB::Exception(DB::ErrorCodes::CANNOT_PARSE_TEXT, "The CIDR has a malformed prefix bits: {}", std::string(cidr_str));
94

95
    return {addr, static_cast<UInt8>(prefix)};
96
}
97

98
inline bool isAddressInRange(const IPAddressVariant & address, const IPAddressCIDR & cidr)
99
{
100
    if (const auto * cidr_v6 = cidr.address.asV6())
101
    {
102
        if (const auto * addr_v6 = address.asV6())
103
            return DB::matchIPv6Subnet(addr_v6, cidr_v6, cidr.prefix);
104
    }
105
    else
106
    {
107
        if (!address.asV6())
108
            return DB::matchIPv4Subnet(address.asV4(), cidr.address.asV4(), cidr.prefix);
109
    }
110
    return false;
111
}
112

113
}
114

115
namespace DB
116
{
117
    class FunctionIsIPAddressContainedIn : public IFunction
118
    {
119
    public:
120
        static constexpr auto name = "isIPAddressInRange";
121
        String getName() const override { return name; }
122
        static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionIsIPAddressContainedIn>(); }
123
        bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
124

125
        ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & /* return_type */, size_t input_rows_count) const override
126
        {
127
            const IColumn * col_addr = arguments[0].column.get();
128
            const IColumn * col_cidr = arguments[1].column.get();
129

130
            if (const auto * col_addr_const = checkAndGetAnyColumnConst(col_addr))
131
            {
132
                if (const auto * col_cidr_const = checkAndGetAnyColumnConst(col_cidr))
133
                    return executeImpl(*col_addr_const, *col_cidr_const, input_rows_count);
134
                else
135
                    return executeImpl(*col_addr_const, *col_cidr, input_rows_count);
136
            }
137
            else
138
            {
139
                if (const auto * col_cidr_const = checkAndGetAnyColumnConst(col_cidr))
140
                    return executeImpl(*col_addr, *col_cidr_const, input_rows_count);
141
                else
142
                    return executeImpl(*col_addr, *col_cidr, input_rows_count);
143
            }
144
        }
145

146
        DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
147
        {
148
            if (arguments.size() != 2)
149
                throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
150
                    "Number of arguments for function {} doesn't match: passed {}, should be 2",
151
                    getName(), arguments.size());
152

153
            const DataTypePtr & addr_type = arguments[0];
154
            const DataTypePtr & prefix_type = arguments[1];
155

156
            if (!isString(addr_type) || !isString(prefix_type))
157
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The arguments of function {} must be String", getName());
158

159
            return std::make_shared<DataTypeUInt8>();
160
        }
161

162
        size_t getNumberOfArguments() const override { return 2; }
163
        bool useDefaultImplementationForNulls() const override { return false; }
164

165
    private:
166
        /// Like checkAndGetColumnConst() but this function doesn't
167
        /// care about the type of data column.
168
        static const ColumnConst * checkAndGetAnyColumnConst(const IColumn * column)
169
        {
170
            if (!column || !isColumnConst(*column))
171
                return nullptr;
172

173
            return assert_cast<const ColumnConst *>(column);
174
        }
175

176
        /// Both columns are constant.
177
        static ColumnPtr executeImpl(
178
            const ColumnConst & col_addr_const,
179
            const ColumnConst & col_cidr_const,
180
            size_t input_rows_count)
181
        {
182
            const auto & col_addr = col_addr_const.getDataColumn();
183
            const auto & col_cidr = col_cidr_const.getDataColumn();
184

185
            const auto addr = IPAddressVariant(col_addr.getDataAt(0).toView());
186
            const auto cidr = parseIPWithCIDR(col_cidr.getDataAt(0).toView());
187

188
            ColumnUInt8::MutablePtr col_res = ColumnUInt8::create(1);
189
            ColumnUInt8::Container & vec_res = col_res->getData();
190

191
            vec_res[0] = isAddressInRange(addr, cidr) ? 1 : 0;
192

193
            return ColumnConst::create(std::move(col_res), input_rows_count);
194
        }
195

196
        /// Address is constant.
197
        static ColumnPtr executeImpl(const ColumnConst & col_addr_const, const IColumn & col_cidr, size_t input_rows_count)
198
        {
199
            const auto & col_addr = col_addr_const.getDataColumn();
200

201
            const auto addr = IPAddressVariant(col_addr.getDataAt(0).toView());
202

203
            ColumnUInt8::MutablePtr col_res = ColumnUInt8::create(input_rows_count);
204
            ColumnUInt8::Container & vec_res = col_res->getData();
205

206
            for (size_t i = 0; i < input_rows_count; ++i)
207
            {
208
                const auto cidr = parseIPWithCIDR(col_cidr.getDataAt(i).toView());
209
                vec_res[i] = isAddressInRange(addr, cidr) ? 1 : 0;
210
            }
211
            return col_res;
212
        }
213

214
        /// CIDR is constant.
215
        static ColumnPtr executeImpl(const IColumn & col_addr, const ColumnConst & col_cidr_const, size_t input_rows_count)
216
        {
217
            const auto & col_cidr = col_cidr_const.getDataColumn();
218

219
            const auto cidr = parseIPWithCIDR(col_cidr.getDataAt(0).toView());
220

221
            ColumnUInt8::MutablePtr col_res = ColumnUInt8::create(input_rows_count);
222
            ColumnUInt8::Container & vec_res = col_res->getData();
223
            for (size_t i = 0; i < input_rows_count; ++i)
224
            {
225
                const auto addr = IPAddressVariant(col_addr.getDataAt(i).toView());
226
                vec_res[i] = isAddressInRange(addr, cidr) ? 1 : 0;
227
            }
228
            return col_res;
229
        }
230

231
        /// Neither are constant.
232
        static ColumnPtr executeImpl(const IColumn & col_addr, const IColumn & col_cidr, size_t input_rows_count)
233
        {
234
            ColumnUInt8::MutablePtr col_res = ColumnUInt8::create(input_rows_count);
235
            ColumnUInt8::Container & vec_res = col_res->getData();
236

237
            for (size_t i = 0; i < input_rows_count; ++i)
238
            {
239
                const auto addr = IPAddressVariant(col_addr.getDataAt(i).toView());
240
                const auto cidr = parseIPWithCIDR(col_cidr.getDataAt(i).toView());
241

242
                vec_res[i] = isAddressInRange(addr, cidr) ? 1 : 0;
243
            }
244

245
            return col_res;
246
        }
247
    };
248

249
    REGISTER_FUNCTION(IsIPAddressContainedIn)
250
    {
251
        factory.registerFunction<FunctionIsIPAddressContainedIn>();
252
    }
253
}
254