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/main | |
parent | 5d42ad8a6ea453c129452cdc23311b6b191ad10a (diff) |
Move CidrBlock to vespa
Diffstat (limited to 'vespajlib/src/main')
-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 |
2 files changed, 232 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; |