diff options
author | Håkon Hallingstad <hakon@yahooinc.com> | 2023-06-12 15:20:41 +0200 |
---|---|---|
committer | Håkon Hallingstad <hakon@yahooinc.com> | 2023-06-12 15:20:41 +0200 |
commit | 85a2298347ea60e6d4f904438995ab8ac8e68dea (patch) | |
tree | da6e6271245113f68fe603e5ac6cfa2177958ffd /vespajlib/src | |
parent | 5d42ad8a6ea453c129452cdc23311b6b191ad10a (diff) |
Move CidrBlock to vespa
Diffstat (limited to 'vespajlib/src')
-rw-r--r-- | vespajlib/src/main/java/ai/vespa/net/CidrBlock.java | 227 | ||||
-rw-r--r-- | vespajlib/src/main/java/ai/vespa/net/package-info.java | 5 | ||||
-rw-r--r-- | vespajlib/src/test/java/ai/vespa/net/CidrBlockTest.java | 150 |
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 |