| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF 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 org.apache.cassandra.auth; |
| |
| import java.net.InetAddress; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.Comparator; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.TreeMap; |
| import java.util.stream.Collectors; |
| |
| import com.google.common.annotations.VisibleForTesting; |
| |
| import org.apache.cassandra.cql3.CIDR; |
| |
| /** |
| * This class implements CIDR Interval tree and the ability to find longest matching CIDR for the given IP. |
| * CIDRs interval tree is a variant of interval tree. Each node contains a CIDR and a value. |
| * A node has left children array and the right children array. |
| * - The left children's CIDRs are either less than the starting IP of parent or overlaps with the parent node. |
| * - The right children's CIDRs are either greater than the ending IP of the parent or overlaps with the parent node. |
| * Note that nodes that overlap with the parent node are included in both left and right children arrays. |
| * |
| * The tree organizes nodes by placing non-overlapping CIDRs at the same level. In general, CIDRs with the same net mask |
| * do not overlap, hence are placed in the same level. CIDRs with different net mask may overlap, hence placed at |
| * different levels in the tree. In addition to this, there is an optimisation to promote a CIDR to an upper level, if |
| * it is not overlapping with any CIDR in the parent level, that means, in such cases a CIDR with different net mask can |
| * co-locate in the same level with other CIDRs. |
| * |
| * Levels closer to the root contains CIDRs with higher net mask value. Net mask value decreases as levels further down |
| * from the root. i.e, Nearer the level to the root, the narrower the CIDR, meaning matching the longer IP prefix. |
| * |
| * Search for Longest matching CIDR for an IP starts at level 0, if not found a match, search continues to the next |
| * level, until it finds a match or reaches leaf nodes without a match. That means search terminates on the first match |
| * closest to the root, i.e, locates narrowest matching CIDR. |
| * |
| * Example: |
| * Assume below CIDRs |
| * "128.10.120.2/10", ==> IP range 128.0.0.0 - 128.63.255.255, netmask 10 |
| * "128.20.120.2/20", ==> IP range 128.20.112.0 - 128.20.127.255, netmask 20 |
| * "0.0.0.0/0", ==> IP range 0.0.0.0 - 255.255.255.255, netmask 0 |
| * "10.1.1.2/10" ==> IP range 10.0.0.0 - 10.63.255.255, netmask 10 |
| * |
| * Resulting interval tree looks like: |
| * |
| * (10.0.0.0 - 10.63.255.255, 10) (128.20.112.0 - 128.20.127.255, 20) |
| * / \ / \ |
| * / (128.0.0.0 - 128.63.255.255, 10) |
| * / / \ |
| * (0.0.0.0 - 255.255.255.255, 0) |
| * |
| * Note that in this example (10.0.0.0 - 10.63.255.255, 10) doesn't have any overlapping CIDR, hence moved up a level as |
| * an optimization |
| */ |
| public class CIDRGroupsMappingIntervalTree<V> implements CIDRGroupsMappingTable<V> |
| { |
| private final IPIntervalTree<V> tree; |
| |
| /** |
| * Build an interval tree for given CIDRs |
| * @param isIPv6 true if this CIDR groups mapping is for IPv6 IPs, false otherwise |
| * @param cidrMappings CIDR to Value mappings |
| */ |
| public CIDRGroupsMappingIntervalTree(boolean isIPv6, Map<CIDR, Set<V>> cidrMappings) |
| { |
| for (CIDR cidr : cidrMappings.keySet()) |
| { |
| if (isIPv6 != cidr.isIPv6()) |
| throw new IllegalArgumentException("Invalid CIDR format, expecting " + getIPTypeString(isIPv6) + |
| ", received " + getIPTypeString(cidr.isIPv6())); |
| } |
| |
| this.tree = IPIntervalTree.build(new ArrayList<>(cidrMappings |
| .entrySet() |
| .stream() |
| .collect(Collectors.groupingBy(p -> p.getKey().getNetMask(), |
| TreeMap::new, |
| Collectors.toList())) |
| .descendingMap() |
| .values())); |
| } |
| |
| /** |
| * Get the longest matching CIDR (i.e, the narrowest match) for given IP |
| * @param ip IP to lookup CIDR group |
| * @return CIDR group name(s) |
| */ |
| public Set<V> lookupLongestMatchForIP(InetAddress ip) |
| { |
| // Valid when empty CIDR groups mappings received, i.e, cidr_groups table is empty |
| if (tree == null) |
| return Collections.emptySet(); |
| |
| return tree.query(ip); |
| } |
| |
| /** |
| * This class represents a node of an IP interval tree. |
| * A node contains a CIDR, value associated with it, left and right children |
| */ |
| static class IPIntervalNode<V> |
| { |
| private final CIDR cidr; |
| private final Set<V> values = new HashSet<>(); |
| private IPIntervalNode<V>[] left; |
| private IPIntervalNode<V>[] right; |
| |
| public IPIntervalNode(CIDR cidr, Set<V> values, IPIntervalNode<V>[] children) |
| { |
| this.cidr = cidr; |
| if (values != null) |
| this.values.addAll(values); |
| updateChildren(children, true, true); |
| } |
| |
| @VisibleForTesting |
| CIDR cidr() |
| { |
| return cidr; |
| } |
| |
| @VisibleForTesting |
| IPIntervalNode<V>[] left() |
| { |
| return left; |
| } |
| |
| @VisibleForTesting |
| IPIntervalNode<V>[] right() |
| { |
| return right; |
| } |
| |
| private void updateLeft(IPIntervalNode<V>[] newValue, boolean shouldUpdate) |
| { |
| if (shouldUpdate) |
| this.left = newValue; |
| } |
| |
| private void updateRight(IPIntervalNode<V>[] newValue, boolean shouldUpdate) |
| { |
| if (shouldUpdate) |
| this.right = newValue; |
| } |
| |
| /** |
| * Split the children array according to the IP range of this node, and link the nodes to either |
| * the left or right children. |
| * @param children array of nodes which are children for this node |
| * @param updateLeft true to update left children, false otherwise |
| * @param updateRight true to update right children, false otherwise |
| */ |
| private void updateChildren(IPIntervalNode<V>[] children, boolean updateLeft, boolean updateRight) |
| { |
| // this is leaf node |
| if (children == null) |
| { |
| updateLeft(null, updateLeft); |
| updateRight(null, updateRight); |
| return; |
| } |
| |
| // Find the node in the children that is the closest to this node. |
| int index = binarySearchNodesIndex(children, this.cidr.getStartIpAddress()); |
| IPIntervalNode<V> closest = children[index]; |
| |
| // Scenario - all children nodes are greater than this node |
| if (index == 0 && CIDR.compareIPs(this.cidr.getEndIpAddress(), closest.cidr.getStartIpAddress()) < 0) |
| { |
| updateLeft(null, updateLeft); |
| updateRight(children, updateRight); |
| } |
| // Scenario - all children nodes are lower than this node |
| else if (index == children.length - 1 && |
| CIDR.compareIPs(this.cidr.getStartIpAddress(), closest.cidr.getEndIpAddress()) > 0) |
| { |
| updateLeft(children, updateLeft); |
| updateRight(null, updateRight); |
| } |
| else // Scenario - part of the children nodes are lower, and the other are greater |
| { |
| // When this node does not overlap with the closest, split the array and |
| // link left and right children correspondingly. |
| if (CIDR.compareIPs(this.cidr.getStartIpAddress(), closest.cidr.getEndIpAddress()) > 0) |
| { |
| // including the closest (node at index) in left |
| updateLeft(Arrays.copyOfRange(children, 0, index + 1), updateLeft); |
| // put the rest in right |
| updateRight(Arrays.copyOfRange(children, index + 1, children.length), updateRight); |
| } |
| else // When the node overlaps, include the closest node in both its left and right children nodes. |
| { |
| // The parent node overlaps with at most 1 interval in the children, because of nature of the CIDR. |
| // Increasing the bit mask by 1, divides the range into halfs. |
| // Note that the node@index is included in both left and right |
| // it is because the current interval partially overlaps with the closest interval |
| // The overlapping interval should always be searched if we cannot find an exact match with current interval. |
| updateLeft(Arrays.copyOfRange(children, 0, index + 1), updateLeft); |
| updateRight(Arrays.copyOfRange(children, index, children.length), updateRight); |
| } |
| } |
| } |
| |
| private void updateLeftIfNull(IPIntervalNode<V>[] children) |
| { |
| if (left != null) |
| return; |
| |
| updateChildren(children, true, false); |
| } |
| |
| private void updateRightIfNull(IPIntervalNode<V>[] children) |
| { |
| if (right != null) |
| return; |
| |
| updateChildren(children, false, true); |
| } |
| |
| /** |
| * Binary search given array of nodes and return index of the closest matching node. |
| * It looks up for the interval that matches exactly or is left (lower) to the ip. |
| * @param nodes array of nodes |
| * @param ip IP address to search |
| * @param <V> data type of the value |
| * @return index of the closest node |
| */ |
| static <V> int binarySearchNodesIndex(IPIntervalNode<V>[] nodes, InetAddress ip) |
| { |
| int start = 0; // inclusive |
| int end = nodes.length; // exclusive |
| |
| while (true) |
| { |
| if (start >= end) |
| { |
| // return the closest |
| return Math.max((end - 1), 0); |
| } |
| |
| int mid = start + (end - start) / 2; |
| IPIntervalNode<V> midNode = nodes[mid]; |
| int cmp = CIDR.compareIPs(ip, midNode.cidr.getStartIpAddress()); |
| |
| if (cmp == 0) // found the node |
| { |
| return mid; |
| } |
| else if (cmp < 0) // Given IP is less than middle node's starting IP, search left side sub array |
| { |
| end = mid; |
| } |
| else // Given IP is >= middle node's starting IP, so compare ending IP |
| { |
| int compEnd = CIDR.compareIPs(ip, midNode.cidr.getEndIpAddress()); |
| // Given IP is >= middle node's starting IP and <= than the ending IP, found the match |
| if (compEnd <= 0) |
| { |
| return mid; |
| } |
| else // IP > middle node's end IP >= given IP, search right side sub array |
| { |
| start = mid + 1; |
| } |
| } |
| } |
| } |
| |
| /** |
| * Binary search given array of nodes to find the closest IP interval to the input IP |
| * @param nodes array of nodes |
| * @param ip IP address to search |
| * @param <V> data type of the value |
| * @return the closest node to the input IP |
| */ |
| static <V> IPIntervalNode<V> binarySearchNodes(IPIntervalNode<V>[] nodes, InetAddress ip) |
| { |
| int index = binarySearchNodesIndex(nodes, ip); |
| return nodes[index]; |
| } |
| |
| /** |
| * Search the tree for a CIDR matching given IP. Uses DFS and stops on first match i.e, finds the closest match |
| * to the root, which is the narrowest matching CIDR |
| * @param root subtree with this node as root |
| * @param ip IP address to search |
| * @param <V> data type of the value |
| * @return value(s) associated with the CIDR matching the given IP, Returns null if no match found |
| */ |
| static <V> Set<V> query(IPIntervalNode<V> root, InetAddress ip) |
| { |
| IPIntervalNode<V> current = root; |
| while (true) // while loop transformed from tail recursion |
| { |
| boolean largerThanStart = CIDR.compareIPs(ip, current.cidr.getStartIpAddress()) >= 0; |
| boolean lessThanEnd = CIDR.compareIPs(ip, current.cidr.getEndIpAddress()) <= 0; |
| if (largerThanStart && lessThanEnd) |
| { |
| return current.values; |
| } |
| else |
| { |
| IPIntervalNode<V>[] candidates = largerThanStart ? current.right : current.left; |
| // the tree is exhausted, and we are unable to find a match |
| if (candidates == null) |
| { |
| return null; |
| } |
| current = binarySearchNodes(candidates, ip); |
| } |
| } |
| } |
| } |
| |
| |
| /** |
| * This class represents an interval tree for CIDRs |
| * @param <V> data type of the value |
| */ |
| static class IPIntervalTree<V> |
| { |
| // References to first level nodes of the tree |
| private final IPIntervalNode<V>[] level0; |
| // depth of the tree |
| private final int depth; |
| |
| private IPIntervalTree(IPIntervalNode<V>[] nodes, int depth) |
| { |
| this.level0 = nodes; |
| this.depth = depth; |
| } |
| |
| @VisibleForTesting |
| int getDepth() |
| { |
| return depth; |
| } |
| |
| /** |
| * Optimize levels by moving non-overlapping CIDRs from lower level to the upper level. Levels are updated in-place |
| * This optimization moves CIDRs closer to the root, hence improves the search to find IP nearer to the root, |
| * i.e, avoiding going depth during the search |
| * @param upperLevel level for CIDRs with higher netmask value |
| * @param lowerLevel level for CIDRs with lower netmask value |
| * @param <V> data type of value |
| */ |
| private static <V> void optimizeLevels(List<Map.Entry<CIDR, V>> upperLevel, List<Map.Entry<CIDR, V>> lowerLevel) |
| { |
| List<Map.Entry<CIDR, V>> newUpper = new ArrayList<>(upperLevel.size() + lowerLevel.size()); |
| newUpper.addAll(upperLevel); |
| |
| List<Map.Entry<CIDR, V>> newLower = new ArrayList<>(lowerLevel.size()); |
| |
| for (int i = 0; i < lowerLevel.size(); i++) |
| { |
| boolean noOverlap = true; |
| for (int j = 0; j < upperLevel.size(); j++) |
| { |
| if (CIDR.overlaps(lowerLevel.get(i).getKey(), upperLevel.get(j).getKey())) |
| { |
| // overlapping node remains in lower level |
| newLower.add(lowerLevel.get(i)); |
| noOverlap = false; |
| break; |
| } |
| } |
| |
| // the node from lower level does not overlap with any node in upper level, so move it up |
| if (noOverlap) |
| { |
| newUpper.add(lowerLevel.get(i)); |
| } |
| } |
| |
| upperLevel.clear(); |
| lowerLevel.clear(); |
| upperLevel.addAll(newUpper); |
| lowerLevel.addAll(newLower); |
| } |
| |
| /** |
| * Optimize by lifting lower level non-overlapping CIDRs to the upper levels. |
| * The method modifies cidrsGroupedByNetMasks in-place. |
| * @param cidrsGroupedByNetMasks CIDRs grouped by netmask, levels in the order of higher netmask to lower netmask value |
| * @param <V> data type of the value |
| */ |
| private static <V> void optimizeAllLevels(List<List<Map.Entry<CIDR, Set<V>>>> cidrsGroupedByNetMasks) |
| { |
| for (int i = 0; i < cidrsGroupedByNetMasks.size(); i++) |
| { |
| List<Map.Entry<CIDR, Set<V>>> current = cidrsGroupedByNetMasks.get(0); |
| for (int j = i + 1; j < cidrsGroupedByNetMasks.size(); j++) |
| { |
| List<Map.Entry<CIDR, Set<V>>> lower = cidrsGroupedByNetMasks.get(j); |
| optimizeLevels(current, lower); |
| } |
| } |
| } |
| |
| /** |
| * Link the nodes between levels. If a node in the lower level has no left or right children, |
| * the nodes from higher levels are lifted and linked as the left or right children. |
| * @param cidrMappings list of CIDR to value mappings |
| * @param result 2 dimension array that stores all (linked) nodes in the tree. result[0] is the first level in the tree. |
| * @param startIndex starting level in the result array, to link it to its children nodes |
| * @param <V> data type of the value |
| */ |
| @SuppressWarnings("unchecked") |
| private static <V> void linkNodes(List<List<Map.Entry<CIDR, Set<V>>>> cidrMappings, |
| IPIntervalNode<V>[][] result, |
| int startIndex) |
| { |
| List<Map.Entry<CIDR, Set<V>>> cidrsAtLevel = cidrMappings.get(startIndex); |
| int next = startIndex + 1; |
| IPIntervalNode<V>[] lowerLevel = next == result.length ? null : result[next]; |
| result[startIndex] = cidrsAtLevel |
| .stream() |
| .map(pair -> { |
| CIDR cidr = pair.getKey(); |
| Set<V> value = pair.getValue(); |
| IPIntervalNode<V> node = new IPIntervalNode<>(cidr, value, lowerLevel); |
| // Cannot link its left or right children to the next level nodes. |
| // Try the nodes in the even lower levels. |
| if (next + 1 < result.length && (node.left == null || node.right == null)) |
| { |
| for (int i = next + 1; i < result.length; i++) |
| { |
| node.updateLeftIfNull(result[i]); |
| node.updateRightIfNull(result[i]); |
| |
| // the node has both left and right children, it is good and no further lifting is needed. |
| if (node.left != null && node.right != null) |
| { |
| break; |
| } |
| } |
| } |
| return node; |
| }) |
| // sort by ipStart. Remember the interval do not overlap |
| .sorted(Comparator.comparing(n -> n.cidr.getStartIpAddress(), CIDR::compareIPs)) |
| .toArray(IPIntervalNode[]::new); |
| } |
| |
| /** |
| * Build an interval tree for given CIDRs |
| * @param cidrsGroupedByNetMasks CIDRs grouped by netmask, levels in the order of higher netmask to lower netmask value |
| * @param <V> data type of the value |
| * @return returns reference to the interval tree, returns null if input is empty |
| */ |
| @SuppressWarnings("unchecked") |
| public static <V> IPIntervalTree<V> build(List<List<Map.Entry<CIDR, Set<V>>>> cidrsGroupedByNetMasks) |
| { |
| if (cidrsGroupedByNetMasks.isEmpty()) |
| return null; |
| |
| optimizeAllLevels(cidrsGroupedByNetMasks); |
| |
| // After optimization, trim empty levels if exists |
| cidrsGroupedByNetMasks.removeIf(List::isEmpty); |
| |
| // Create an array for each level |
| IPIntervalNode<V>[][] result = new IPIntervalNode[cidrsGroupedByNetMasks.size()][]; |
| // build the tree bottom up, i.e, add CIDRs in the order of smaller netmask value to higher netmask value |
| for (int i = cidrsGroupedByNetMasks.size() - 1; i >= 0; i--) |
| { |
| linkNodes(cidrsGroupedByNetMasks, result, i); |
| } |
| return new IPIntervalTree<>(result[0], cidrsGroupedByNetMasks.size()); |
| } |
| |
| /** |
| * Search interval tree for the longest matching CIDR for the given IP |
| * @param ip IP address to search |
| * @return Value(s) associated with matching CIDR |
| */ |
| public Set<V> query(InetAddress ip) |
| { |
| IPIntervalNode<V> closest = IPIntervalNode.binarySearchNodes(level0, ip); |
| return IPIntervalNode.query(closest, ip); |
| } |
| } |
| } |
