Move CidrBlock to vespa

3 files changed, 382 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/ai/vespa/net/CidrBlock.java b/vespajlib/src/main/java/ai/vespa/net/CidrBlock.java
new file mode 100644
index 00000000000..533e5fcde5e
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/net/CidrBlock.java
@@ -0,0 +1,227 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.net;
+
+import com.google.common.net.InetAddresses;
+
+import java.io.UncheckedIOException;
+import java.math.BigInteger;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.Objects;
+
+/**
+ * Represents a single IPv4 or IPv6 CIDR block.
+ *
+ * @author valerijf
+ */
+public class CidrBlock {
+
+    private final BigInteger ipAddressBits;
+    private final int prefixLength;
+    private final int addressLength;
+
+    /** Creates a CIDR block that only contains the provided IP address (/32 if IPv4, /128 if IPv6) */
+    public CidrBlock(InetAddress inetAddress) {
+        this(inetAddress, 8 * inetAddress.getAddress().length);
+    }
+
+    public CidrBlock(InetAddress inetAddress, int prefixLength) {
+        this(inetAddress.getAddress(), prefixLength);
+    }
+
+    public CidrBlock(byte[] address, int prefixLength) {
+        if (prefixLength < 0) throw new IllegalArgumentException(
+                "Prefix size cannot be negative, but was " + prefixLength);
+
+        this.prefixLength = prefixLength;
+        this.addressLength = 8 * address.length;
+        if (prefixLength > addressLength) throw new IllegalArgumentException(
+                String.format("Prefix size (%s) cannot be longer than address length (%s)", prefixLength, addressLength));
+
+        this.ipAddressBits = inetAddressToBits(address, prefixLength);
+    }
+
+    /** For internal use only, does not validate */
+    private CidrBlock(BigInteger ipAddressBits, int prefixLength, int addressLength) {
+        this.ipAddressBits = ipAddressBits;
+        this.prefixLength = prefixLength;
+        this.addressLength = addressLength;
+    }
+
+    /** @return The first IP address in this CIDR block */
+    public InetAddress getInetAddress() {
+        return bitsToInetAddress(ipAddressBits, addressLength);
+    }
+
+    /** @return the number of bits in the network mask */
+    public int prefixLength() {
+        return prefixLength;
+    }
+
+    public boolean isIpv6() {
+        return addressLength == 128;
+    }
+
+    /** Returns a copy of this resized to the given newPrefixLength */
+    public CidrBlock resize(int newPrefixLength) {
+        return new CidrBlock(ipAddressBits, newPrefixLength, addressLength);
+    }
+
+    public CidrBlock clearLeastSignificantBits(int bits) {
+        return new CidrBlock(ipAddressBits.shiftRight(bits).shiftLeft(bits), prefixLength, addressLength);
+    }
+
+    /** @return a copy of this CIDR block with the host identifier bits cleared */
+    public CidrBlock clearHostIdentifier() {
+        return clearLeastSignificantBits(addressLength - prefixLength);
+    }
+
+    /** Return the byte at the given offset.  0 refers to the most significant byte of the address. */
+    public int getByte(int byteOffset) {
+        return ipAddressBits.shiftRight(addressLength - 8 * (byteOffset + 1)).and(BigInteger.valueOf(0xFF)).intValueExact();
+    }
+
+    /** Set the byte at the given offset to 'n'.  0 refers to the most significant byte of the address. */
+    public CidrBlock setByte(int byteOffset, int n) {
+        if (n < 0 || n > 0xFF) throw new IllegalArgumentException("Byte value must be between 0 and 255, but was " + n);
+        int byteDiff = n - getByte(byteOffset);
+        return addByteRaw(byteOffset, byteDiff);
+    }
+
+    /** Add 'n' to the byte at the given offset, truncating overflow bits.  0 refers to the most significant byte of the address. */
+    public CidrBlock addByte(int byteOffset, int n) {
+        int oldByte = getByte(byteOffset);
+        int newByte = 0xFF & (oldByte + n);
+        return addByteRaw(byteOffset, newByte - oldByte);
+    }
+
+    private CidrBlock addByteRaw(int byteOffset, int n) {
+        BigInteger bit = ipAddressBits.add(BigInteger.valueOf(n).shiftLeft(addressLength - 8 * (byteOffset + 1)));
+        return new CidrBlock(bit, prefixLength, addressLength);
+    }
+
+    public boolean overlapsWith(CidrBlock other) {
+        if (this.isIpv6() != other.isIpv6()) return false;
+
+        int ignoreLastNBits = addressLength - Math.min(this.prefixLength(), other.prefixLength());
+        return this.ipAddressBits.shiftRight(ignoreLastNBits).equals(other.ipAddressBits.shiftRight(ignoreLastNBits));
+    }
+
+    /** @return the .arpa address for this CIDR block, does not include bit outside the prefix */
+    public String getDomainName() {
+        StringBuilder recordPtr = new StringBuilder(75);
+        int segmentWidth = isIpv6() ? 4 : 8;
+
+        int start = addressLength - prefixLength - (segmentWidth - (prefixLength % segmentWidth)) % segmentWidth;
+        for (int i = start; i < addressLength; i += segmentWidth) {
+            int segment = ipAddressBits.shiftRight(i)
+                    .and(BigInteger.ONE.shiftLeft(segmentWidth).subtract(BigInteger.ONE))
+                    .intValueExact();
+
+            recordPtr.append(isIpv6() ? Integer.toHexString(segment) : segment).append(".");
+        }
+
+        return recordPtr.append(isIpv6() ? "ip6" : "in-addr").append(".arpa.").toString();
+    }
+
+    /** @return iterable over all CIDR blocks of the same prefix size, from the current one and up */
+    public Iterable<CidrBlock> iterableCidrs() {
+        return () -> new Iterator<>() {
+            private final BigInteger increment = BigInteger.ONE.shiftLeft(addressLength - prefixLength);
+            private final BigInteger maxValue = BigInteger.ONE.shiftLeft(addressLength).subtract(increment);
+            private BigInteger current = ipAddressBits;
+
+            public boolean hasNext() {
+                return current.compareTo(maxValue) < 0;
+            }
+
+            public CidrBlock next() {
+                if (!hasNext()) throw new NoSuchElementException();
+                CidrBlock cidrBlock = new CidrBlock(current, prefixLength, addressLength);
+                current = current.add(increment);
+                return cidrBlock;
+            }
+        };
+    }
+
+    public Iterable<InetAddress> iterableIps() {
+        return () -> new Iterator<>() {
+            private final BigInteger maxValue = ipAddressBits.or(BigInteger.ONE.shiftLeft(addressLength - prefixLength).subtract(BigInteger.ONE));
+            private BigInteger current = ipAddressBits;
+
+            public boolean hasNext() {
+                return current.compareTo(maxValue) <= 0;
+            }
+
+            public InetAddress next() {
+                if (!hasNext()) throw new NoSuchElementException();
+                InetAddress inetAddress = bitsToInetAddress(current, addressLength);
+                current = current.add(BigInteger.ONE);
+                return inetAddress;
+            }
+        };
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        CidrBlock cidrBlock = (CidrBlock) o;
+        return prefixLength == cidrBlock.prefixLength &&
+               Objects.equals(ipAddressBits, cidrBlock.ipAddressBits);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(ipAddressBits, prefixLength);
+    }
+
+    @Override
+    public String toString() {
+        return asString();
+    }
+
+    public String asString() {
+        return InetAddresses.toAddrString(getInetAddress()) + "/" + prefixLength;
+    }
+
+    public static CidrBlock fromString(String cidr) {
+        String[] cidrParts = cidr.split("/");
+        if (cidrParts.length != 2)
+            throw new IllegalArgumentException("Invalid CIDR block, expected format to be " +
+                    "'<ip address>/<prefix size>', but was '" + cidr + "'");
+
+        InetAddress inetAddress = InetAddresses.forString(cidrParts[0]);
+        int prefixSize = Integer.parseInt(cidrParts[1]);
+
+        return new CidrBlock(inetAddress, prefixSize);
+    }
+
+    private static BigInteger inetAddressToBits(byte[] address, int prefix) {
+        BigInteger bit = BigInteger.ZERO;
+        for (byte b : address)
+            bit = bit.shiftLeft(8).add(BigInteger.valueOf(b & 0xFF));
+        return bit;
+    }
+
+    private static InetAddress bitsToInetAddress(BigInteger ipAddressBits, int addressLength) {
+        try {
+            byte[] addr = ipAddressBits.toByteArray();
+            int addressBytes = addressLength / 8;
+            if (addr.length != addressBytes) {
+                byte[] temp = new byte[addressBytes];
+                System.arraycopy(
+                        addr, Math.max(addr.length - addressBytes, 0),
+                        temp, Math.max(addressBytes - addr.length, 0), Math.min(addr.length, addressBytes));
+                addr = temp;
+            }
+            return InetAddress.getByAddress(addr);
+        } catch (UnknownHostException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+}
diff --git a/vespajlib/src/main/java/ai/vespa/net/package-info.java b/vespajlib/src/main/java/ai/vespa/net/package-info.java
new file mode 100644
index 00000000000..5d5bb613870
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/net/package-info.java
@@ -0,0 +1,5 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+@ExportPackage
+package ai.vespa.net;
+
+import com.yahoo.osgi.annotation.ExportPackage;
diff --git a/vespajlib/src/test/java/ai/vespa/net/CidrBlockTest.java b/vespajlib/src/test/java/ai/vespa/net/CidrBlockTest.java
new file mode 100644
index 00000000000..2194ca2de71
--- /dev/null
+++ b/vespajlib/src/test/java/ai/vespa/net/CidrBlockTest.java
@@ -0,0 +1,150 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.net;
+
+import com.google.common.net.InetAddresses;
+import ai.vespa.net.CidrBlock;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+
+/**
+ * @author valerijf
+ */
+public class CidrBlockTest {
+
+    @Test(expected = IllegalArgumentException.class)
+    public void negative_prefix_size_fails() {
+        new CidrBlock(InetAddresses.forString("10.0.0.1"), -1);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void too_large_ipv4_prefix_size_fails() {
+        new CidrBlock(InetAddresses.forString("10.0.0.1"), 35);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void too_large_ipv6_prefix_size_fails() {
+        new CidrBlock(InetAddresses.forString("::1"), 130);
+    }
+
+    @Test
+    public void bits_after_prefix_are_ignored() {
+        withHostIdentifierCleared("10.3.0.0/16", "10.3.12.180/16", "10.3.128.3/16");
+        withHostIdentifierCleared("1234:5600::/25", "1234:5678::123/25", "1234:5600::725/25");
+    }
+
+    private void withHostIdentifierCleared(String canonical, String... aliases) {
+        CidrBlock canonicalBlock = CidrBlock.fromString(canonical);
+        assertEquals(canonical, canonicalBlock.toString());
+        for (String alias : aliases) {
+            CidrBlock aliasBlock = CidrBlock.fromString(alias).clearHostIdentifier();
+            assertEquals(canonicalBlock, aliasBlock);
+            assertEquals(canonical, aliasBlock.toString());
+        }
+    }
+
+    @Test
+    public void parse_from_string_test() {
+        assertEquals(new CidrBlock(InetAddresses.forString("10.0.0.1"), 12), CidrBlock.fromString("10.0.0.1/12"));
+        assertEquals(new CidrBlock(InetAddresses.forString("1234:5678::1"), 64), CidrBlock.fromString("1234:5678:0000::1/64"));
+    }
+
+    @Test
+    public void ipv4_overlap_test() {
+        CidrBlock base = CidrBlock.fromString("10.3.0.0/16");
+
+        assertTrue(base.overlapsWith(CidrBlock.fromString("10.3.0.0/16")));
+        assertTrue(base.overlapsWith(CidrBlock.fromString("10.0.0.0/8")));
+        assertTrue(base.overlapsWith(CidrBlock.fromString("10.3.128.0/17")));
+        assertFalse(base.overlapsWith(CidrBlock.fromString("10.4.0.0/16")));
+        assertFalse(base.overlapsWith(CidrBlock.fromString("11.0.0.0/8")));
+    }
+
+    @Test
+    public void ipv6_overlap_test() {
+        CidrBlock base = CidrBlock.fromString("1234:5678:abcd::/48");
+
+        assertTrue(base.overlapsWith(CidrBlock.fromString("1234:5678:abcd::/48")));
+        assertTrue(base.overlapsWith(CidrBlock.fromString("1234:5678::/32")));
+        assertTrue(base.overlapsWith(CidrBlock.fromString("1234:5678:abcd:8000::/49")));
+        assertFalse(base.overlapsWith(CidrBlock.fromString("1234:5678:abce::/48")));
+        assertFalse(base.overlapsWith(CidrBlock.fromString("1234:5679::/32")));
+    }
+
+    @Test
+    public void domain_name_test() {
+        assertEquals("3.10.in-addr.arpa.", CidrBlock.fromString("10.3.0.0/16").getDomainName());
+        assertEquals("128.192.in-addr.arpa.", CidrBlock.fromString("192.128.0.0/9").getDomainName());
+
+        assertEquals("d.c.b.a.8.7.6.5.4.3.2.1.ip6.arpa.", CidrBlock.fromString("1234:5678:abcd::/48").getDomainName());
+        assertEquals("8.c.b.a.8.7.6.5.4.3.2.1.ip6.arpa.", CidrBlock.fromString("1234:5678:abc8::/45").getDomainName());
+    }
+
+    @Test
+    public void iterableCidrs() {
+        CidrBlock superBlock = CidrBlock.fromString("10.12.14.0/24");
+        assertEquals(List.of("10.12.14.200/29", "10.12.14.208/29", "10.12.14.216/29", "10.12.14.224/29", "10.12.14.232/29", "10.12.14.240/29", "10.12.14.248/29"),
+                StreamSupport.stream(CidrBlock.fromString("10.12.14.200/29").iterableCidrs().spliterator(), false)
+                        .takeWhile(superBlock::overlapsWith)
+                        .map(CidrBlock::asString)
+                        .collect(Collectors.toList()));
+
+        assertEquals(StreamSupport.stream(superBlock.iterableIps().spliterator(), false)
+                        .skip(24)
+                        .map(ip -> InetAddresses.toAddrString(ip) + "/32")
+                        .collect(Collectors.toList()),
+                StreamSupport.stream(CidrBlock.fromString("10.12.14.24/32").iterableCidrs().spliterator(), false)
+                        .takeWhile(superBlock::overlapsWith)
+                        .map(CidrBlock::asString)
+                        .collect(Collectors.toList()));
+    }
+
+    @Test
+    public void iterableIps() {
+        assertEquals(List.of("10.12.14.24", "10.12.14.25", "10.12.14.26", "10.12.14.27", "10.12.14.28", "10.12.14.29", "10.12.14.30", "10.12.14.31"),
+                StreamSupport.stream(CidrBlock.fromString("10.12.14.24/29").iterableIps().spliterator(), false)
+                        .map(InetAddresses::toAddrString)
+                        .collect(Collectors.toList()));
+
+        assertEquals(List.of("10.12.14.24"),
+                StreamSupport.stream(CidrBlock.fromString("10.12.14.24/32").iterableIps().spliterator(), false)
+                        .map(InetAddresses::toAddrString)
+                        .collect(Collectors.toList()));
+    }
+
+    @Test
+    public void testByteAccessors() {
+        CidrBlock ipv4Block = CidrBlock.fromString("11.22.33.44/24");
+        assertEquals(11, ipv4Block.getByte(0));
+        assertEquals(22, ipv4Block.getByte(1));
+        assertEquals(33, ipv4Block.getByte(2));
+        assertEquals(44, ipv4Block.getByte(3));
+        assertEquals(24, ipv4Block.prefixLength());
+
+        assertEquals("12.22.33.44/24", ipv4Block.addByte(0, 1).asString());
+        assertEquals("11.22.33.45/24", ipv4Block.addByte(3, 1).asString());
+        assertEquals("55.22.33.44/24", ipv4Block.setByte(0, 55).asString());
+        assertEquals("11.22.33.55/24", ipv4Block.setByte(3, 55).asString());
+
+        CidrBlock ipv6Block = CidrBlock.fromString("de63:aca:dcdc:e2a8:1fb4:542:b80d:f7c3/24");
+        assertEquals(0xde, ipv6Block.getByte(0));
+        assertEquals(0x63, ipv6Block.getByte(1));
+        assertEquals(0x0a, ipv6Block.getByte(2));
+        assertEquals(0xc3, ipv6Block.getByte(15));
+        assertEquals(24, ipv6Block.prefixLength());
+
+        assertEquals("df63:aca:dcdc:e2a8:1fb4:542:b80d:f7c3/24", ipv6Block.addByte(0, 1).asString());
+        assertEquals("ab63:aca:dcdc:e2a8:1fb4:542:b80d:f7c3/24", ipv6Block.setByte(0, 0xab).asString());
+        assertEquals("de63:aca:dcdc:e2a8:1fb4:542:b80d:f7c4/24", ipv6Block.addByte(15, 1).asString());
+        assertEquals("de63:aca:dcdc:e2a8:1fb4:542:b80d:f7fe/24", ipv6Block.setByte(15, 0xfe).asString());
+    }
+}
\ No newline at end of file