From 0040ed7cd9232535af4afb2c0318db109263c39a Mon Sep 17 00:00:00 2001
From: takonomura <takonomura@users.noreply.github.com>
Date: Fri, 7 Mar 2025 17:42:01 +0900
Subject: [PATCH] pkg/net: add ip arithmetic functions

This change introduces two functions to adding numerical offsets to IP
addresses or CIDR networks.

Signed-off-by: takonomura <takonomura@users.noreply.github.com>
---
 pkg/net/ip.go              | 61 ++++++++++++++++++++++++++++++++++++++
 pkg/net/pkg.go             | 26 ++++++++++++++++
 pkg/net/testdata/gen.txtar | 14 +++++++++
 3 files changed, 101 insertions(+)

diff --git a/pkg/net/ip.go b/pkg/net/ip.go
index 89c587be167..e6871f1b046 100644
--- a/pkg/net/ip.go
+++ b/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
+}
diff --git a/pkg/net/pkg.go b/pkg/net/pkg.go
index 568236636be..7abb12990e2 100644
--- a/pkg/net/pkg.go
+++ b/pkg/net/pkg.go
@@ -237,6 +237,32 @@ var p = &pkg.Package{
 				c.Ret, c.Err = IPString(ip)
 			}
 		},
+	}, {
+		Name: "AddIPOffset",
+		Params: []pkg.Param{
+			{Kind: adt.TopKind},
+			{Kind: adt.IntKind},
+		},
+		Result: adt.StringKind,
+		Func: func(c *pkg.CallCtxt) {
+			ip, offset := c.Value(0), c.BigInt(1)
+			if c.Do() {
+				c.Ret, c.Err = AddIPOffset(ip, offset)
+			}
+		},
+	}, {
+		Name: "AddIPCIDROffset",
+		Params: []pkg.Param{
+			{Kind: adt.TopKind},
+			{Kind: adt.IntKind},
+		},
+		Result: adt.StringKind,
+		Func: func(c *pkg.CallCtxt) {
+			ip, offset := c.Value(0), c.BigInt(1)
+			if c.Do() {
+				c.Ret, c.Err = AddIPCIDROffset(ip, offset)
+			}
+		},
 	}, {
 		Name: "PathEscape",
 		Params: []pkg.Param{
diff --git a/pkg/net/testdata/gen.txtar b/pkg/net/testdata/gen.txtar
index 28023c57a9b..f3fbd28417c 100644
--- a/pkg/net/testdata/gen.txtar
+++ b/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"