Take suffixes into account in case of format_ipv4_mask

ClibMouse · Aug 26, 2023 · 105a0e0 · 105a0e0
1 parent 4579e5c
commit 105a0e0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/Parsers/Kusto/KustoFunctions/KQLIPFunctions.cpp b/src/Parsers/Kusto/KustoFunctions/KQLIPFunctions.cpp
@@ -264,10 +264,12 @@ bool FormatIpv4Mask::convertImpl(String & out, IParser::Pos & pos)
     const auto mask = getOptionalArgument(function_name, pos);
     const auto calculated_mask = mask ? *mask : "32";
     out = std::format(
-        "if(empty({1} as formatted_ip_{2}) or not {0} between 0 and 32, '', concat(formatted_ip_{2}, '/', toString({0})))",
+        "if(empty({1} as formatted_ip_{2}) or position(toTypeName({0}), 'Int') = 0 or not {0} between 0 and 32, '', "
+        "concat(formatted_ip_{2}, '/', toString(toInt64(min2({0}, ifNull({3} as suffix_{2}, 32))))))",
         calculated_mask,
         kqlCallToExpression("format_ipv4", {ip_address, calculated_mask}, pos.max_depth),
-        generateUniqueIdentifier());
+        generateUniqueIdentifier(),
+        kqlCallToExpression("ipv4_netmask_suffix", {"tostring("+ip_address+")"}, pos.max_depth));
     return true;
 }
 }
diff --git a/src/Parsers/tests/KQL/gtest_KQL_IP.cpp b/src/Parsers/tests/KQL/gtest_KQL_IP.cpp
@@ -16,11 +16,11 @@ INSTANTIATE_TEST_SUITE_P(ParserKQLQuery_IP, ParserRegexTest,
         },
         {
             "print format_ipv4_mask(A)",
-            "SELECT if\\(empty\\(ifNull\\(if\\(\\(\\(\\(toUInt32OrNull\\(toString\\(A\\)\\) AS param_as_uint32_\\d+\\) IS NOT NULL\\) AND \\(toTypeName\\(A\\) = 'String'\\)\\) OR \\(32 < 0\\) OR \\(\\(ifNull\\(param_as_uint32_\\d+, multiIf\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) = 1, IPv4StringToNumOrNull\\(tokens_\\d+\\[1\\]\\) AS ip_\\d+, \\(length\\(tokens_\\d+\\) = 2\\) AND \\(ip_\\d+ IS NOT NULL\\) AND \\(\\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NOT NULL\\), IPv4CIDRToRange\\(assumeNotNull\\(ip_\\d+\\), assumeNotNull\\(mask_\\d+\\)\\).1, NULL\\)\\) AS ip_as_number_\\d+\\) IS NULL\\), NULL, IPv4NumToString\\(bitAnd\\(ip_as_number_\\d+, bitNot\\(toUInt32\\(intExp2\\(32 - 32\\) - 1\\)\\)\\)\\)\\), ''\\) AS formatted_ip_\\d+\\) OR \\(NOT \\(\\(32 >= 0\\) AND \\(32 <= 32\\)\\)\\), '', concat\\(formatted_ip_\\d+, '/', toString\\(32\\)\\)\\)"
+            "SELECT if\\(empty\\(ifNull\\(if\\(\\(\\(\\(toUInt32OrNull\\(toString\\(A\\)\\) AS param_as_uint32_\\d+\\) IS NOT NULL\\) AND \\(toTypeName\\(A\\) = 'String'\\)\\) OR \\(32 < 0\\) OR \\(\\(ifNull\\(param_as_uint32_\\d+, multiIf\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) = 1, IPv4StringToNumOrNull\\(tokens_\\d+\\[1\\]\\) AS ip_\\d+, \\(length\\(tokens_\\d+\\) = 2\\) AND \\(ip_\\d+ IS NOT NULL\\) AND \\(\\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NOT NULL\\), IPv4CIDRToRange\\(assumeNotNull\\(ip_\\d+\\), assumeNotNull\\(mask_\\d+\\)\\)\\.1, NULL\\)\\) AS ip_as_number_\\d+\\) IS NULL\\), NULL, IPv4NumToString\\(bitAnd\\(ip_as_number_\\d+, bitNot\\(toUInt32\\(intExp2\\(32 - 32\\) - 1\\)\\)\\)\\)\\), ''\\) AS formatted_ip_\\d+\\) OR \\(position\\(toTypeName\\(32\\), 'Int'\\) = 0\\) OR \\(NOT \\(\\(32 >= 0\\) AND \\(32 <= 32\\)\\)\\), '', concat\\(formatted_ip_\\d+, '/', toString\\(toInt64\\(min2\\(32, ifNull\\(multiIf\\(\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) > 2\\) OR \\(NOT isIPv4String\\(tokens_\\d+\\[1\\]\\)\\), NULL, length\\(tokens_\\d+\\) = 1, 32, \\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NULL, NULL, toUInt8\\(min2\\(mask_\\d+, 32\\)\\)\\) AS suffix_\\d+, 32\\)\\)\\)\\)\\)\\)"
         },
         {
             "print format_ipv4_mask(A, B)",
-            "SELECT if\\(empty\\(ifNull\\(if\\(\\(\\(\\(toUInt32OrNull\\(toString\\(A\\)\\) AS param_as_uint32_\\d+\\) IS NOT NULL\\) AND \\(toTypeName\\(A\\) = 'String'\\)\\) OR \\(B < 0\\) OR \\(\\(ifNull\\(param_as_uint32_\\d+, multiIf\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) = 1, IPv4StringToNumOrNull\\(tokens_\\d+\\[1\\]\\) AS ip_\\d+, \\(length\\(tokens_\\d+\\) = 2\\) AND \\(ip_\\d+ IS NOT NULL\\) AND \\(\\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NOT NULL\\), IPv4CIDRToRange\\(assumeNotNull\\(ip_\\d+\\), assumeNotNull\\(mask_\\d+\\)\\).1, NULL\\)\\) AS ip_as_number_\\d+\\) IS NULL\\), NULL, IPv4NumToString\\(bitAnd\\(ip_as_number_\\d+, bitNot\\(toUInt32\\(intExp2\\(32 - B\\) - 1\\)\\)\\)\\)\\), ''\\) AS formatted_ip_\\d+\\) OR \\(NOT \\(\\(B >= 0\\) AND \\(B <= 32\\)\\)\\), '', concat\\(formatted_ip_\\d+, '/', toString\\(B\\)\\)\\)"
+            "SELECT if\\(empty\\(ifNull\\(if\\(\\(\\(\\(toUInt32OrNull\\(toString\\(A\\)\\) AS param_as_uint32_\\d+\\) IS NOT NULL\\) AND \\(toTypeName\\(A\\) = 'String'\\)\\) OR \\(B < 0\\) OR \\(\\(ifNull\\(param_as_uint32_\\d+, multiIf\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) = 1, IPv4StringToNumOrNull\\(tokens_\\d+\\[1\\]\\) AS ip_\\d+, \\(length\\(tokens_\\d+\\) = 2\\) AND \\(ip_\\d+ IS NOT NULL\\) AND \\(\\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NOT NULL\\), IPv4CIDRToRange\\(assumeNotNull\\(ip_\\d+\\), assumeNotNull\\(mask_\\d+\\)\\)\\.1, NULL\\)\\) AS ip_as_number_\\d+\\) IS NULL\\), NULL, IPv4NumToString\\(bitAnd\\(ip_as_number_\\d+, bitNot\\(toUInt32\\(intExp2\\(32 - B\\) - 1\\)\\)\\)\\)\\), ''\\) AS formatted_ip_\\d+\\) OR \\(position\\(toTypeName\\(B\\), 'Int'\\) = 0\\) OR \\(NOT \\(\\(B >= 0\\) AND \\(B <= 32\\)\\)\\), '', concat\\(formatted_ip_\\d+, '/', toString\\(toInt64\\(min2\\(B, ifNull\\(multiIf\\(\\(length\\(splitByChar\\('/', ifNull\\(toString\\(A\\), ''\\)\\) AS tokens_\\d+\\) > 2\\) OR \\(NOT isIPv4String\\(tokens_\\d+\\[1\\]\\)\\), NULL, length\\(tokens_\\d+\\) = 1, 32, \\(toUInt8OrNull\\(tokens_\\d+\\[-1\\]\\) AS mask_\\d+\\) IS NULL, NULL, toUInt8\\(min2\\(mask_\\d+, 32\\)\\)\\) AS suffix_\\d+, 32\\)\\)\\)\\)\\)\\)"
         },
         {
             "print ipv4_compare(A, B)",