From bebebac02e3ad59a3686cd6c9a221807c2c60068 Mon Sep 17 00:00:00 2001 From: mikhail-khludnev Date: Thu, 17 Oct 2024 19:10:40 +0300 Subject: [PATCH] max cap test --- .../search/SearchIpFieldTermsTest.java | 155 +++++++++++++++--- 1 file changed, 135 insertions(+), 20 deletions(-) diff --git a/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTest.java b/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTest.java index d437d353290f5..84104707c8b88 100644 --- a/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTest.java +++ b/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTest.java @@ -8,54 +8,169 @@ package org.opensearch.search; +import org.opensearch.action.bulk.BulkRequestBuilder; import org.opensearch.action.search.SearchResponse; +import org.opensearch.common.network.InetAddresses; import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.test.OpenSearchSingleNodeTestCase; import java.io.IOException; +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; import static org.hamcrest.Matchers.equalTo; import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE; public class SearchIpFieldTermsTest extends OpenSearchSingleNodeTestCase { + public static final boolean IPv4_ONLY = true; static String defaultIndexName = "test"; - public void testOneAddr() throws Exception { + public void testMassive() throws Exception { XContentBuilder xcb = createMapping(); client().admin().indices().prepareCreate(defaultIndexName).setMapping(xcb).get(); ensureGreen(); - client().prepareIndex(defaultIndexName) - .setSource("{\"addr\": \"192.168.0.1\"}", MediaTypeRegistry.JSON) - .setRefreshPolicy(IMMEDIATE) + BulkRequestBuilder bulkRequestBuilder = client().prepareBulk(); + + int cidrs = 0; + int ips = 0; + List toQuery = new ArrayList<>(); + for (int i = 0; ips <= 1024; i++) { + final String ip; + final int prefix; + if (IPv4_ONLY) { + ip = generateRandomIPv4(); + prefix = 8 + random().nextInt(24); // CIDR prefix for IPv4 + } else { + ip = generateRandomIPv6(); + prefix = 32 + random().nextInt(97); // CIDR prefix for IPv6 + } + + bulkRequestBuilder.add(client().prepareIndex(defaultIndexName).setSource(Map.of("addr", ip))); + + final String termToQuery; + if (cidrs < 1024 - 1 && random().nextBoolean()) { + termToQuery = ip + "/" + prefix; + cidrs++; + } else { + termToQuery = ip; + ips++; + } + toQuery.add(termToQuery); + } + int addMatches = 0; + for (int i = 0; i < atLeast(100); i++) { + final String ip; + if (IPv4_ONLY) { + ip = generateRandomIPv4(); + } else { + ip = generateRandomIPv6(); + } + bulkRequestBuilder.add(client().prepareIndex(defaultIndexName).setSource(Map.of("addr", ip))); + boolean match = false; + for (String termQ : toQuery) { + boolean isCidr = termQ.contains("/"); + if ((isCidr && isIPInCIDR(ip, termQ)) || (!isCidr && termQ.equals(ip))) { + match = true; + break; + } + } + if (match) { + addMatches++; + } else { + break; // single mismatch is enough. + } + } + + bulkRequestBuilder.setRefreshPolicy(IMMEDIATE) .get(); SearchResponse result = client().prepareSearch(defaultIndexName) - .setQuery(QueryBuilders.termsQuery("addr","192.168.0.1" )) + .setQuery(QueryBuilders.termsQuery("addr", toQuery)) .get(); - assertThat(result.getHits().getTotalHits().value, equalTo(1L)); + assertThat(result.getHits().getTotalHits().value, equalTo(0L + cidrs + ips + addMatches)); + } + + + // Converts an IP string (either IPv4 or IPv6) to a byte array + private static byte[] ipToBytes(String ip) { + InetAddress inetAddress = InetAddresses.forString(ip); + return inetAddress.getAddress(); + } + + // Checks if an IP is within a given CIDR (works for both IPv4 and IPv6) + private static boolean isIPInCIDR(String ip, String cidr) throws UnknownHostException { + String[] cidrParts = cidr.split("/"); + String cidrIp = cidrParts[0]; + int prefixLength = Integer.parseInt(cidrParts[1]); + + byte[] ipBytes = ipToBytes(ip); + byte[] cidrIpBytes = ipToBytes(cidrIp); + + // Calculate how many full bytes and how many bits are in the mask + int fullBytes = prefixLength / 8; + int extraBits = prefixLength % 8; + + // Compare full bytes + for (int i = 0; i < fullBytes; i++) { + if (ipBytes[i] != cidrIpBytes[i]) { + return false; + } + } + + // Compare extra bits (if any) + if (extraBits > 0) { + int mask = 0xFF << (8 - extraBits); + if ((ipBytes[fullBytes] & mask) != (cidrIpBytes[fullBytes] & mask)) { + return false; + } + } + + return true; + } + + // Generate a random IPv4 address + private static String generateRandomIPv4() { + return String.format("%d.%d.%d.%d", random().nextInt(256), + random().nextInt(256), + random().nextInt(256), + random().nextInt(256)); + } + + // Generate a random IPv6 address + private static String generateRandomIPv6() { + StringBuilder ipv6 = new StringBuilder(); + for (int i = 0; i < 8; i++) { + ipv6.append(Integer.toHexString(random().nextInt(0xFFFF + 1))); + if (i < 7) { + ipv6.append(":"); + } + } + return ipv6.toString(); } private XContentBuilder createMapping() throws IOException { XContentBuilder xcb = XContentFactory.jsonBuilder() .startObject() .startObject("properties") - .startObject("addr") - .field("type", "ip") - .startObject("fields") - .startObject("idx") - .field("type","ip") - .field("doc_values",false) - .endObject() - .startObject("dv") - .field("type","ip") - .field("index",false) - .endObject() - .endObject() - .endObject() + .startObject("addr") + .field("type", "ip") + .startObject("fields") + .startObject("idx") + .field("type", "ip") + .field("doc_values", false) + .endObject() + .startObject("dv") + .field("type", "ip") + .field("index", false) + .endObject() + .endObject() + .endObject() .endObject() .endObject(); return xcb;