Binary search based IpSubnetFilter (#10492)

Motivation:

`IpSubnetFilter` uses Binary Search for IP Address search which is fast if we have large set of IP addresses to filter.

Modification:

Added `IpSubnetFilter` which takes `IpSubnetFilterRule` for filtering.

Result:
Faster IP address filter.

handler/src/main/java/io/netty/handler/ipfilter/IpSubnetFilter.java

@@ -0,0 +1,226 @@
+/*
+ * Copyright 2020 The Netty Project
+ *
+ * The Netty Project licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package io.netty.handler.ipfilter;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandler.Sharable;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.util.internal.ObjectUtil;
+
+import java.net.Inet4Address;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * <p>
+ * This class allows one to filter new {@link Channel}s based on the
+ * {@link IpSubnetFilter}s passed to its constructor. If no rules are provided, all connections
+ * will be accepted since {@code acceptIfNotFound} is {@code true} by default.
+ * </p>
+ *
+ * <p>
+ * If you would like to explicitly take action on rejected {@link Channel}s, you should override
+ * {@link AbstractRemoteAddressFilter#channelRejected(ChannelHandlerContext, SocketAddress)}.
+ * </p>
+ *
+ * <p>
+ *     Few Points to keep in mind:
+ *     <ol>
+ *         <li> Since {@link IpSubnetFilter} uses Binary search algorithm, it's a good
+ *         idea to insert IP addresses in incremental order. </li>
+ *         <li> Remove any over-lapping CIDR.  </li>
+ *     </ol>
+ * </p>
+ *
+ */
+@Sharable
+public class IpSubnetFilter extends AbstractRemoteAddressFilter<InetSocketAddress> {
+
+    private final boolean acceptIfNotFound;
+    private final List<IpSubnetFilterRule> ipv4Rules;
+    private final List<IpSubnetFilterRule> ipv6Rules;
+    private final IpFilterRuleType ipFilterRuleTypeIPv4;
+    private final IpFilterRuleType ipFilterRuleTypeIPv6;
+
+    /**
+     * <p> Create new {@link IpSubnetFilter} Instance with specified {@link IpSubnetFilterRule} as array. </p>
+     * <p> {@code acceptIfNotFound} is set to {@code true}. </p>
+     *
+     * @param rules {@link IpSubnetFilterRule} as an array
+     */
+    public IpSubnetFilter(IpSubnetFilterRule... rules) {
+        this(true, Arrays.asList(ObjectUtil.checkNotNull(rules, "rules")));
+    }
+
+    /**
+     * <p> Create new {@link IpSubnetFilter} Instance with specified {@link IpSubnetFilterRule} as array
+     * and specify if we'll accept a connection if we don't find it in the rule(s). </p>
+     *
+     * @param acceptIfNotFound {@code true} if we'll accept connection if not found in rule(s).
+     * @param rules            {@link IpSubnetFilterRule} as an array
+     */
+    public IpSubnetFilter(boolean acceptIfNotFound, IpSubnetFilterRule... rules) {
+        this(acceptIfNotFound, Arrays.asList(ObjectUtil.checkNotNull(rules, "rules")));
+    }
+
+    /**
+     * <p> Create new {@link IpSubnetFilter} Instance with specified {@link IpSubnetFilterRule} as {@link List}. </p>
+     * <p> {@code acceptIfNotFound} is set to {@code true}. </p>
+     *
+     * @param rules {@link IpSubnetFilterRule} as a {@link List}
+     */
+    public IpSubnetFilter(List<IpSubnetFilterRule> rules) {
+        this(true, rules);
+    }
+
+    /**
+     * <p> Create new {@link IpSubnetFilter} Instance with specified {@link IpSubnetFilterRule} as {@link List}
+     * and specify if we'll accept a connection if we don't find it in the rule(s). </p>
+     *
+     * @param acceptIfNotFound {@code true} if we'll accept connection if not found in rule(s).
+     * @param rules            {@link IpSubnetFilterRule} as a {@link List}
+     */
+    public IpSubnetFilter(boolean acceptIfNotFound, List<IpSubnetFilterRule> rules) {
+        ObjectUtil.checkNotNull(rules, "rules");
+        this.acceptIfNotFound = acceptIfNotFound;
+
+        int numAcceptIPv4 = 0;
+        int numRejectIPv4 = 0;
+        int numAcceptIPv6 = 0;
+        int numRejectIPv6 = 0;
+
+        List<IpSubnetFilterRule> unsortedIPv4Rules = new ArrayList<IpSubnetFilterRule>();
+        List<IpSubnetFilterRule> unsortedIPv6Rules = new ArrayList<IpSubnetFilterRule>();
+
+        // Iterate over rules and check for `null` rule.
+        for (IpSubnetFilterRule ipSubnetFilterRule : rules) {
+            ObjectUtil.checkNotNull(ipSubnetFilterRule, "rule");
+
+            if (ipSubnetFilterRule.getFilterRule() instanceof IpSubnetFilterRule.Ip4SubnetFilterRule) {
+                unsortedIPv4Rules.add(ipSubnetFilterRule);
+
+                if (ipSubnetFilterRule.ruleType() == IpFilterRuleType.ACCEPT) {
+                    numAcceptIPv4++;
+                } else {
+                    numRejectIPv4++;
+                }
+            } else {
+                unsortedIPv6Rules.add(ipSubnetFilterRule);
+
+                if (ipSubnetFilterRule.ruleType() == IpFilterRuleType.ACCEPT) {
+                    numAcceptIPv6++;
+                } else {
+                    numRejectIPv6++;
+                }
+            }
+        }
+
+        /*
+         * If Number of ACCEPT rule is 0 and number of REJECT rules is more than 0,
+         * then all rules are of "REJECT" type.
+         *
+         * In this case, we'll set `ipFilterRuleTypeIPv4` to `IpFilterRuleType.REJECT`.
+         *
+         * If Number of ACCEPT rules are more than 0 and number of REJECT rules is 0,
+         * then all rules are of "ACCEPT" type.
+         *
+         * In this case, we'll set `ipFilterRuleTypeIPv4` to `IpFilterRuleType.ACCEPT`.
+         */
+        if (numAcceptIPv4 == 0 && numRejectIPv4 > 0) {
+            ipFilterRuleTypeIPv4 = IpFilterRuleType.REJECT;
+        } else if (numAcceptIPv4 > 0 && numRejectIPv4 == 0) {
+            ipFilterRuleTypeIPv4 = IpFilterRuleType.ACCEPT;
+        } else {
+            ipFilterRuleTypeIPv4 = null;
+        }
+
+        if (numAcceptIPv6 == 0 && numRejectIPv6 > 0) {
+            ipFilterRuleTypeIPv6 = IpFilterRuleType.REJECT;
+        } else if (numAcceptIPv6 > 0 && numRejectIPv6 == 0) {
+            ipFilterRuleTypeIPv6 = IpFilterRuleType.ACCEPT;
+        } else {
+            ipFilterRuleTypeIPv6 = null;
+        }
+
+        this.ipv4Rules = sortAndFilter(unsortedIPv4Rules);
+        this.ipv6Rules = sortAndFilter(unsortedIPv6Rules);
+    }
+
+    @Override
+    protected boolean accept(ChannelHandlerContext ctx, InetSocketAddress remoteAddress) {
+        if (remoteAddress.getAddress() instanceof Inet4Address) {
+            int indexOf = Collections.binarySearch(ipv4Rules, remoteAddress, IpSubnetFilterRuleComparator.INSTANCE);
+            if (indexOf >= 0) {
+                if (ipFilterRuleTypeIPv4 == null) {
+                    return ipv4Rules.get(indexOf).ruleType() == IpFilterRuleType.ACCEPT;
+                } else {
+                    return ipFilterRuleTypeIPv4 == IpFilterRuleType.ACCEPT;
+                }
+            }
+        } else {
+            int indexOf = Collections.binarySearch(ipv6Rules, remoteAddress, IpSubnetFilterRuleComparator.INSTANCE);
+            if (indexOf >= 0) {
+                if (ipFilterRuleTypeIPv6 == null) {
+                    return ipv6Rules.get(indexOf).ruleType() == IpFilterRuleType.ACCEPT;
+                } else {
+                    return ipFilterRuleTypeIPv6 == IpFilterRuleType.ACCEPT;
+                }
+            }
+        }
+
+        return acceptIfNotFound;
+    }
+
+    /**
+     * <ol>
+     *     <li> Sort the list </li>
+     *     <li> Remove over-lapping subnet </li>
+     *     <li> Sort the list again </li>
+     * </ol>
+     */
+    @SuppressWarnings("ConstantConditions")
+    private static List<IpSubnetFilterRule> sortAndFilter(List<IpSubnetFilterRule> rules) {
+        Collections.sort(rules);
+        Iterator<IpSubnetFilterRule> iterator = rules.iterator();
+        List<IpSubnetFilterRule> toKeep = new ArrayList<IpSubnetFilterRule>();
+
+        IpSubnetFilterRule parentRule = iterator.hasNext() ? iterator.next() : null;
+        if (parentRule != null) {
+            toKeep.add(parentRule);
+        }
+
+        while (iterator.hasNext()) {
+
+            // Grab a potential child rule.
+            IpSubnetFilterRule childRule = iterator.next();
+
+            // If parentRule matches childRule, then there's no need to keep the child rule.
+            // Otherwise, the rules are distinct and we need both.
+            if (!parentRule.matches(new InetSocketAddress(childRule.getIpAddress(), 1))) {
+                toKeep.add(childRule);
+                // Then we'll keep the child rule around as the parent for the next round.
+                parentRule = childRule;
+            }
+        }
+
+        return toKeep;
+    }
+}

handler/src/main/java/io/netty/handler/ipfilter/IpSubnetFilterRule.java

@@ -17,6 +17,7 @@ package io.netty.handler.ipfilter;
 
 import static java.util.Objects.requireNonNull;
 
+import io.netty.util.NetUtil;
 import io.netty.util.internal.SocketUtils;
 
 import java.math.BigInteger;
@@ -30,12 +31,14 @@ import java.net.UnknownHostException;
  * Use this class to create rules for {@link RuleBasedIpFilter} that group IP addresses into subnets.
  * Supports both, IPv4 and IPv6.
  */
-public final class IpSubnetFilterRule implements IpFilterRule {
+public final class IpSubnetFilterRule implements IpFilterRule, Comparable<IpSubnetFilterRule> {
 
     private final IpFilterRule filterRule;
+    private final String ipAddress;
 
     public IpSubnetFilterRule(String ipAddress, int cidrPrefix, IpFilterRuleType ruleType) {
         try {
+            this.ipAddress = ipAddress;
             filterRule = selectFilterRule(SocketUtils.addressByName(ipAddress), cidrPrefix, ruleType);
         } catch (UnknownHostException e) {
             throw new IllegalArgumentException("ipAddress", e);
@@ -43,6 +46,7 @@ public final class IpSubnetFilterRule implements IpFilterRule {
     }
 
     public IpSubnetFilterRule(InetAddress ipAddress, int cidrPrefix, IpFilterRuleType ruleType) {
+        this.ipAddress = ipAddress.getHostAddress();
         filterRule = selectFilterRule(ipAddress, cidrPrefix, ruleType);
     }
 
@@ -69,7 +73,59 @@ public final class IpSubnetFilterRule implements IpFilterRule {
         return filterRule.ruleType();
     }
 
-    private static final class Ip4SubnetFilterRule implements IpFilterRule {
+    /**
+     * Get IP Address of this rule
+     */
+    String getIpAddress() {
+        return ipAddress;
+    }
+
+    /**
+     * {@link Ip4SubnetFilterRule} or {@link Ip6SubnetFilterRule}
+     */
+    IpFilterRule getFilterRule() {
+        return filterRule;
+    }
+
+    @Override
+    public int compareTo(IpSubnetFilterRule ipSubnetFilterRule) {
+        if (filterRule instanceof Ip4SubnetFilterRule) {
+            return compareInt(((Ip4SubnetFilterRule) filterRule).networkAddress,
+                    ((Ip4SubnetFilterRule) ipSubnetFilterRule.filterRule).networkAddress);
+        } else {
+            return ((Ip6SubnetFilterRule) filterRule).networkAddress
+                    .compareTo(((Ip6SubnetFilterRule) ipSubnetFilterRule.filterRule).networkAddress);
+        }
+    }
+
+    /**
+     * It'll compare IP address with {@link Ip4SubnetFilterRule#networkAddress} or
+     * {@link Ip6SubnetFilterRule#networkAddress}.
+     *
+     * @param inetSocketAddress {@link InetSocketAddress} to match
+     * @return 0 if IP Address match else difference index.
+     */
+    int compareTo(InetSocketAddress inetSocketAddress) {
+        if (filterRule instanceof Ip4SubnetFilterRule) {
+            Ip4SubnetFilterRule ip4SubnetFilterRule = (Ip4SubnetFilterRule) filterRule;
+            return compareInt(ip4SubnetFilterRule.networkAddress, NetUtil.ipv4AddressToInt((Inet4Address)
+                    inetSocketAddress.getAddress()) & ip4SubnetFilterRule.subnetMask);
+        } else {
+            Ip6SubnetFilterRule ip6SubnetFilterRule = (Ip6SubnetFilterRule) filterRule;
+            return ip6SubnetFilterRule.networkAddress
+                    .compareTo(Ip6SubnetFilterRule.ipToInt((Inet6Address) inetSocketAddress.getAddress())
+                            .and(ip6SubnetFilterRule.networkAddress));
+        }
+    }
+
+    /**
+     * Equivalent to {@link Integer#compare(int, int)}
+     */
+    private static int compareInt(int x, int y) {
+        return (x < y) ? -1 : ((x == y) ? 0 : 1);
+    }
+
+    static final class Ip4SubnetFilterRule implements IpFilterRule {
 
         private final int networkAddress;
         private final int subnetMask;
@@ -77,12 +133,12 @@ public final class IpSubnetFilterRule implements IpFilterRule {
 
         private Ip4SubnetFilterRule(Inet4Address ipAddress, int cidrPrefix, IpFilterRuleType ruleType) {
             if (cidrPrefix < 0 || cidrPrefix > 32) {
-                throw new IllegalArgumentException(String.format("IPv4 requires the subnet prefix to be in range of " +
+                throw new IllegalArgumentException(String.format("IPv4 requires the subnet prefix to be in range of "
-                                                                    "[0,32]. The prefix was: %d", cidrPrefix));
+                        + "[0,32]. The prefix was: %d", cidrPrefix));
             }
 
             subnetMask = prefixToSubnetMask(cidrPrefix);
-            networkAddress = ipToInt(ipAddress) & subnetMask;
+            networkAddress = NetUtil.ipv4AddressToInt(ipAddress) & subnetMask;
             this.ruleType = ruleType;
         }
 
@@ -90,7 +146,7 @@ public final class IpSubnetFilterRule implements IpFilterRule {
         public boolean matches(InetSocketAddress remoteAddress) {
             final InetAddress inetAddress = remoteAddress.getAddress();
             if (inetAddress instanceof Inet4Address) {
-                int ipAddress = ipToInt((Inet4Address) inetAddress);
+                int ipAddress = NetUtil.ipv4AddressToInt((Inet4Address) inetAddress);
                 return (ipAddress & subnetMask) == networkAddress;
             }
             return false;
@@ -101,18 +157,8 @@ public final class IpSubnetFilterRule implements IpFilterRule {
             return ruleType;
         }
 
-        private static int ipToInt(Inet4Address ipAddress) {
-            byte[] octets = ipAddress.getAddress();
-            assert octets.length == 4;
-
-            return (octets[0] & 0xff) << 24 |
-                   (octets[1] & 0xff) << 16 |
-                   (octets[2] & 0xff) << 8 |
-                    octets[3] & 0xff;
-        }
-
         private static int prefixToSubnetMask(int cidrPrefix) {
-            /**
+            /*
              * Perform the shift on a long and downcast it to int afterwards.
              * This is necessary to handle a cidrPrefix of zero correctly.
              * The left shift operator on an int only uses the five least
@@ -126,7 +172,7 @@ public final class IpSubnetFilterRule implements IpFilterRule {
         }
     }
 
-    private static final class Ip6SubnetFilterRule implements IpFilterRule {
+    static final class Ip6SubnetFilterRule implements IpFilterRule {
 
         private static final BigInteger MINUS_ONE = BigInteger.valueOf(-1);
 
@@ -136,8 +182,8 @@ public final class IpSubnetFilterRule implements IpFilterRule {
 
         private Ip6SubnetFilterRule(Inet6Address ipAddress, int cidrPrefix, IpFilterRuleType ruleType) {
             if (cidrPrefix < 0 || cidrPrefix > 128) {
-                throw new IllegalArgumentException(String.format("IPv6 requires the subnet prefix to be in range of " +
+                throw new IllegalArgumentException(String.format("IPv6 requires the subnet prefix to be in range of "
-                                                                    "[0,128]. The prefix was: %d", cidrPrefix));
+                        + "[0,128]. The prefix was: %d", cidrPrefix));
             }
 
             subnetMask = prefixToSubnetMask(cidrPrefix);

handler/src/main/java/io/netty/handler/ipfilter/IpSubnetFilterRuleComparator.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2020 The Netty Project
+ *
+ * The Netty Project licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package io.netty.handler.ipfilter;
+
+import java.net.InetSocketAddress;
+import java.util.Comparator;
+
+/**
+ * This comparator is only used for searching.
+ */
+final class IpSubnetFilterRuleComparator implements Comparator<Object> {
+
+    static final IpSubnetFilterRuleComparator INSTANCE = new IpSubnetFilterRuleComparator();
+
+    private IpSubnetFilterRuleComparator() {
+        // Prevent outside initialization
+    }
+
+    @Override
+    public int compare(Object o1, Object o2) {
+        return ((IpSubnetFilterRule) o1).compareTo((InetSocketAddress) o2);
+    }
+}

handler/src/test/java/io/netty/handler/ipfilter/IpSubnetFilterTest.java

@@ -27,6 +27,9 @@ import org.junit.Test;
 
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
 
 public class IpSubnetFilterTest {
 
@@ -141,6 +144,57 @@ public class IpSubnetFilterTest {
         Assert.assertTrue(ch4.isActive());
     }
 
+    @Test
+    public void testBinarySearch() {
+        List<IpSubnetFilterRule> ipSubnetFilterRuleList = new ArrayList<IpSubnetFilterRule>();
+        ipSubnetFilterRuleList.add(buildRejectIP("1.2.3.4", 32));
+        ipSubnetFilterRuleList.add(buildRejectIP("1.1.1.1", 8));
+        ipSubnetFilterRuleList.add(buildRejectIP("200.200.200.200", 32));
+        ipSubnetFilterRuleList.add(buildRejectIP("108.0.0.0", 4));
+        ipSubnetFilterRuleList.add(buildRejectIP("10.10.10.10", 8));
+        ipSubnetFilterRuleList.add(buildRejectIP("2001:db8:abcd:0000::", 52));
+
+        // 1.0.0.0/8
+        EmbeddedChannel ch1 = newEmbeddedInetChannel("1.1.1.1", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertFalse(ch1.isActive());
+        Assert.assertTrue(ch1.close().isSuccess());
+
+        // Nothing applies here
+        EmbeddedChannel ch2 = newEmbeddedInetChannel("2.2.2.2", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertTrue(ch2.isActive());
+        Assert.assertTrue(ch2.close().isSuccess());
+
+        // 108.0.0.0/4
+        EmbeddedChannel ch3 = newEmbeddedInetChannel("97.100.100.100", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertFalse(ch3.isActive());
+        Assert.assertTrue(ch3.close().isSuccess());
+
+        // 200.200.200.200/32
+        EmbeddedChannel ch4 = newEmbeddedInetChannel("200.200.200.200", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertFalse(ch4.isActive());
+        Assert.assertTrue(ch4.close().isSuccess());
+
+        // Nothing applies here
+        EmbeddedChannel ch5 = newEmbeddedInetChannel("127.0.0.1", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertTrue(ch5.isActive());
+        Assert.assertTrue(ch5.close().isSuccess());
+
+        // 10.0.0.0/8
+        EmbeddedChannel ch6 = newEmbeddedInetChannel("10.1.1.2", new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertFalse(ch6.isActive());
+        Assert.assertTrue(ch6.close().isSuccess());
+
+        //2001:db8:abcd:0000::/52
+        EmbeddedChannel ch7 = newEmbeddedInetChannel("2001:db8:abcd:1000::",
+                new IpSubnetFilter(ipSubnetFilterRuleList));
+        Assert.assertFalse(ch7.isActive());
+        Assert.assertTrue(ch7.close().isSuccess());
+    }
+
+    private static IpSubnetFilterRule buildRejectIP(String ipAddress, int mask) {
+        return new IpSubnetFilterRule(ipAddress, mask, IpFilterRuleType.REJECT);
+    }
+
     private static EmbeddedChannel newEmbeddedInetChannel(final String ipAddress, ChannelHandler... handlers) {
         return new EmbeddedChannel(handlers) {
             @Override