go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/netutil/cidr.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

/*
NOTE: this file is largely based on the work done here: https://github.com/apparentlymart/go-cidr

The original pre-amble is as follows:

Package cidr is a collection of assorted utilities for computing
network and host addresses within network ranges.

It expects a CIDR-type address structure where addresses are divided into
some number of prefix bits representing the network and then the remaining
suffix bits represent the host.

For example, it can help to calculate addresses for sub-networks of a
parent network, or to calculate host addresses within a particular prefix.

At present this package is prioritizing simplicity of implementation and
de-prioritizing speed and memory usage. Thus caution is advised before
using this package in performance-critical applications or hot codepaths.
Patches to improve the speed and memory usage may be accepted as long as
they do not result in a significant increase in code complexity.

*/

package netutil

import (
	"fmt"
	"math/big"
	"net"
)

// ParseCIDR wraps net.ParseCIDR and returns an augmented net.IPNet
// that can do range "math" e.g. `.Next()`.
func ParseCIDR(cidr string) (net.IP, *IPNet, error) {
	ip, m, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, nil, err
	}
	if m == nil {
		return nil, nil, nil
	}
	return ip, &IPNet{m}, nil
}

// IPNet is a wrapper for a `net.IPNet` with additional helper functions.
type IPNet struct {
	*net.IPNet
}

// AddressCount returns the number of distinct host addresses within the given
// CIDR range.
//
// Since the result is a uint64, this function returns meaningful information
// only for IPv4 ranges and IPv6 ranges with a prefix size of at least 65.
func (i IPNet) AddressCount() uint64 {
	prefixLen, bits := i.Mask.Size()
	return 1 << (uint64(bits) - uint64(prefixLen))
}

// Next produces the next, non-overlapping mask for this mask.
//
// An example:
//
//	_, r0, _ := netutil.ParseCIDR("192.168.0.0/24")
//	r1 := r1.Next()
//	fmt.Println(r1.String()) // prints "192.168.1.0/24"
//
// It will return false if the ip range is not incrementable, that is
// if the combination of address and size are the last combination representable.
func (i IPNet) Next() (*IPNet, bool) {
	_, currentLast := i.Range()
	prefixLen, _ := i.Mask.Size()
	mask := net.CIDRMask(prefixLen, 8*len(currentLast))
	currentSubnet := &IPNet{&net.IPNet{IP: currentLast.Mask(mask), Mask: mask}}
	_, last := currentSubnet.Range()
	last = incrementAddress(last)
	if last.Equal(net.IPv4zero) || last.Equal(net.IPv6zero) {
		return nil, false
	}
	next := &IPNet{&net.IPNet{IP: last.Mask(mask), Mask: mask}}
	return next, true
}

// Range returns the first and last IP address of the IPNet.
func (i IPNet) Range() (net.IP, net.IP) {
	firstIP := i.IP
	prefixLen, bits := i.Mask.Size()
	if prefixLen == bits {
		lastIP := make([]byte, len(firstIP))
		copy(lastIP, firstIP)
		return firstIP, lastIP
	}
	firstIPInt, bits, _ := ipToBigInt(firstIP)
	hostLen := uint(bits) - uint(prefixLen)
	lastIPInt := big.NewInt(1)
	lastIPInt.Lsh(lastIPInt, hostLen)
	lastIPInt.Sub(lastIPInt, big.NewInt(1))
	lastIPInt.Or(lastIPInt, firstIPInt)
	return firstIP, bigIntToIP(lastIPInt, bits)
}

// Overlaps returns true if this subnet overlaps with a given subnet.
func (i IPNet) Overlaps(other *IPNet) bool {
	first, last := i.Range()
	otherFirst, otherLast := other.Range()
	if i.Contains(otherFirst) || i.Contains(otherLast) {
		return true
	}
	if other.Contains(first) || other.Contains(last) {
		return true
	}
	return false
}

// incrementAddress increases the IP by one this returns a new []byte for the IP
func incrementAddress(ip net.IP) net.IP {
	ip = maybeResizeV4(ip)
	incIP := make([]byte, len(ip))
	copy(incIP, ip)
	for j := len(incIP) - 1; j >= 0; j-- {
		incIP[j]++
		if incIP[j] > 0 {
			break
		}
	}
	return incIP
}

func ipToBigInt(ip net.IP) (*big.Int, int, error) {
	val := new(big.Int)
	val.SetBytes([]byte(ip))
	if len(ip) == net.IPv4len {
		return val, 32, nil
	} else if len(ip) == net.IPv6len {
		return val, 128, nil
	} else {
		return nil, 0, fmt.Errorf("cidr; unsupported address length %d", len(ip))
	}
}

func bigIntToIP(ipInt *big.Int, bits int) net.IP {
	ipBytes := ipInt.Bytes()
	ret := make([]byte, bits>>3)
	for i := 1; i <= len(ipBytes); i++ {
		ret[len(ret)-i] = ipBytes[len(ipBytes)-i]
	}
	return net.IP(ret)
}

func maybeResizeV4(ip net.IP) net.IP {
	if v4 := ip.To4(); v4 != nil {
		return v4
	}
	return ip
}