pkg/net: add ip arithmetic functions

pkg/net/ip.go

@@ -16,7 +16,9 @@
 package net
 
 import (
+	"errors"
 	"fmt"
+	"math/big"
 	"net/netip"
 
 	"cuelang.org/go/cue"
@@ -242,3 +244,62 @@ func IPString(ip cue.Value) (string, error) {
 	}
 	return ipdata.String(), nil
 }
+
+func netIPAddOffset(addr netip.Addr, offset *big.Int) (netip.Addr, error) {
+	i := big.NewInt(0).SetBytes(addr.AsSlice())
+	i = i.Add(i, offset)
+
+	if i.Sign() < 0 {
+		return netip.Addr{}, errors.New("underflowed")
+	}
+
+	b := i.Bytes()
+	size := addr.BitLen() / 8
+
+	if len(b) > size {
+		return netip.Addr{}, errors.New("overflowed")
+	}
+
+	if len(b) < size {
+		b = append(make([]byte, size-len(b)), b...)
+	}
+	addr, _ = netip.AddrFromSlice(b)
+	return addr, nil
+}
+
+// AddIPOffset adds a numerical offset to a given IP address.
+// The address can be provided as a string, byte array, or CIDR subnet notation.
+// It returns the resulting IP address or CIDR subnet notation as a string.
+func AddIPOffset(ip cue.Value, offset *big.Int) (string, error) {
+	prefix, err := netGetIPCIDR(ip)
+	if err == nil {
+		addr, err := netIPAddOffset(prefix.Addr(), offset)
+		if err != nil {
+			return "", err
+		}
+		return netip.PrefixFrom(addr, prefix.Bits()).String(), nil
+	}
+	ipdata := netGetIP(ip)
+	if !ipdata.IsValid() {
+		return "", fmt.Errorf("invalid IP %q", ip)
+	}
+	addr, err := netIPAddOffset(ipdata, offset)
+	if err != nil {
+		return "", err
+	}
+	return addr.String(), nil
+}
+
+// AddIPCIDROffset adds a numerical offset to a given CIDR subnet.
+func AddIPCIDROffset(ip cue.Value, offset *big.Int) (string, error) {
+	prefix, err := netGetIPCIDR(ip)
+	if err != nil {
+		return "", err
+	}
+	offset.Lsh(offset, (uint)(prefix.Addr().BitLen()-prefix.Bits()))
+	addr, err := netIPAddOffset(prefix.Addr(), offset)
+	if err != nil {
+		return "", err
+	}
+	return netip.PrefixFrom(addr, prefix.Bits()).String(), nil
+}

pkg/net/testdata/gen.txtar

@@ -37,6 +37,13 @@ t31: net.URL & "https://foo.com/bar"
 t32: net.AbsURL & "/foo/bar"
 t33: net.AbsURL & "https://foo.com/bar"
 t34: net.AbsURL & "%"
+t35: net.AddIPOffset("127.0.0.1", 1)
+t36: net.AddIPOffset("127.0.0.1/8", 2)
+t37: net.AddIPOffset("2001:db8::", 1)
+t38: net.AddIPOffset("2001:db8::/64", 2)
+t39: net.AddIPOffset("::ffff:127.0.0.1", 1)
+t40: net.AddIPCIDROffset("192.168.0.0/23", 1)
+t41: net.AddIPCIDROffset("2001:db8:1111:2222::/64", 1)
 -- out/net --
 Errors:
 t25: invalid value "2001:db8::1234567" (does not satisfy net.IPv6):
@@ -100,3 +107,10 @@ t31: "https://foo.com/bar"
 t32: _|_ // t32: invalid value "/foo/bar" (does not satisfy net.AbsURL): t32: error in call to net.AbsURL: URL is not absolute
 t33: "https://foo.com/bar"
 t34: _|_ // t34: invalid value "%" (does not satisfy net.AbsURL): t34: error in call to net.AbsURL: parse "%": invalid URL escape "%"
+t35: "127.0.0.2"
+t36: "127.0.0.3/8"
+t37: "2001:db8::1"
+t38: "2001:db8::2/64"
+t39: "::ffff:127.0.0.2"
+t40: "192.168.2.0/23"
+t41: "2001:db8:1111:2223::/64"