Add support for `net` package IP and CIDR parsing together with CIDR functions which are working similar to Terraform IP network functions

2 files changed, 238 insertions, 0 deletions

diff --git a/funcs/net.go b/funcs/net.go
index b110dbdf..a2d1397c 100644
--- a/funcs/net.go
+++ b/funcs/net.go
@@ -2,10 +2,13 @@ package funcs
 
 import (
 	"context"
+	"math/big"
 	stdnet "net"
 
+	"github.com/apparentlymart/go-cidr/cidr"
 	"github.com/hairyhenderson/gomplate/v3/conv"
 	"github.com/hairyhenderson/gomplate/v3/net"
+	"github.com/pkg/errors"
 	"inet.af/netaddr"
 )
 
@@ -80,3 +83,136 @@ func (f NetFuncs) ParseIPPrefix(ipprefix interface{}) (netaddr.IPPrefix, error)
 func (f NetFuncs) ParseIPRange(iprange interface{}) (netaddr.IPRange, error) {
 	return netaddr.ParseIPRange(conv.ToString(iprange))
 }
+
+// StdParseIP -
+func (f NetFuncs) StdParseIP(prefix interface{}) (stdnet.IP, error) {
+	ip := stdnet.ParseIP(conv.ToString(prefix))
+	if ip == nil {
+		return nil, errors.Errorf("invalid IP address")
+	}
+	return ip, nil
+}
+
+func (f NetFuncs) stdParseCIDR(prefix interface{}) (*stdnet.IPNet, error) {
+	if n, ok := prefix.(*stdnet.IPNet); ok {
+		return n, nil
+	}
+
+	_, network, err := stdnet.ParseCIDR(conv.ToString(prefix))
+	return network, err
+}
+
+// StdParseCIDR -
+func (f NetFuncs) StdParseCIDR(prefix interface{}) (*stdnet.IPNet, error) {
+	return f.stdParseCIDR(prefix)
+}
+
+// CidrHost -
+func (f NetFuncs) CidrHost(hostnum interface{}, prefix interface{}) (*stdnet.IP, error) {
+	network, err := f.stdParseCIDR(prefix)
+	if err != nil {
+		return nil, err
+	}
+
+	ip, err := cidr.HostBig(network, big.NewInt(conv.ToInt64(hostnum)))
+	return &ip, err
+}
+
+// CidrNetmask -
+func (f NetFuncs) CidrNetmask(prefix interface{}) (*stdnet.IP, error) {
+	network, err := f.stdParseCIDR(prefix)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(network.IP) != stdnet.IPv4len {
+		return nil, errors.Errorf("only IPv4 networks are supported")
+	}
+
+	netmask := stdnet.IP(network.Mask)
+	return &netmask, nil
+}
+
+// CidrSubnets -
+func (f NetFuncs) CidrSubnets(newbits interface{}, prefix interface{}) ([]*stdnet.IPNet, error) {
+	network, err := f.stdParseCIDR(prefix)
+	if err != nil {
+		return nil, err
+	}
+
+	nBits := conv.ToInt(newbits)
+	if nBits < 1 {
+		return nil, errors.Errorf("must extend prefix by at least one bit")
+	}
+
+	maxNetNum := int64(1 << uint64(nBits))
+	retValues := make([]*stdnet.IPNet, maxNetNum)
+	for i := int64(0); i < maxNetNum; i++ {
+		subnet, err := cidr.SubnetBig(network, nBits, big.NewInt(i))
+		if err != nil {
+			return nil, err
+		}
+		retValues[i] = subnet
+	}
+
+	return retValues, nil
+}
+
+// CidrSubnetSizes -
+func (f NetFuncs) CidrSubnetSizes(args ...interface{}) ([]*stdnet.IPNet, error) {
+	if len(args) < 2 {
+		return nil, errors.Errorf("wrong number of args: want 2 or more, got %d", len(args))
+	}
+
+	network, err := f.stdParseCIDR(args[len(args)-1])
+	if err != nil {
+		return nil, err
+	}
+	newbits := conv.ToInts(args[:len(args)-1]...)
+
+	startPrefixLen, _ := network.Mask.Size()
+	firstLength := newbits[0]
+
+	firstLength += startPrefixLen
+	retValues := make([]*stdnet.IPNet, len(newbits))
+
+	current, _ := cidr.PreviousSubnet(network, firstLength)
+
+	for i, length := range newbits {
+		if length < 1 {
+			return nil, errors.Errorf("must extend prefix by at least one bit")
+		}
+		// For portability with 32-bit systems where the subnet number
+		// will be a 32-bit int, we only allow extension of 32 bits in
+		// one call even if we're running on a 64-bit machine.
+		// (Of course, this is significant only for IPv6.)
+		if length > 32 {
+			return nil, errors.Errorf("may not extend prefix by more than 32 bits")
+		}
+
+		length += startPrefixLen
+		if length > (len(network.IP) * 8) {
+			protocol := "IP"
+			switch len(network.IP) {
+			case stdnet.IPv4len:
+				protocol = "IPv4"
+			case stdnet.IPv6len:
+				protocol = "IPv6"
+			}
+			return nil, errors.Errorf("would extend prefix to %d bits, which is too long for an %s address", length, protocol)
+		}
+
+		next, rollover := cidr.NextSubnet(current, length)
+		if rollover || !network.Contains(next.IP) {
+			// If we run out of suffix bits in the base CIDR prefix then
+			// NextSubnet will start incrementing the prefix bits, which
+			// we don't allow because it would then allocate addresses
+			// outside of the caller's given prefix.
+			return nil, errors.Errorf("not enough remaining address space for a subnet with a prefix of %d bits after %s", length, current.String())
+		}
+		current = next
+		retValues[i] = current
+	}
+
+	return retValues, nil
+}
diff --git a/funcs/net_test.go b/funcs/net_test.go
index 0c556516..eaa4ea65 100644
--- a/funcs/net_test.go
+++ b/funcs/net_test.go
@@ -2,6 +2,7 @@ package funcs
 
 import (
 	"context"
+	stdnet "net"
 	"strconv"
 	"testing"
 
@@ -69,3 +70,104 @@ func TestParseIPRange(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "192.168.0.2-192.168.23.255", iprange.String())
 }
+
+func TestStdParseIP(t *testing.T) {
+	n := NetFuncs{}
+	ip, err := n.StdParseIP("not an IP")
+	assert.Nil(t, ip)
+	assert.Error(t, err)
+
+	ip, err = n.StdParseIP("10.12.113.12")
+	assert.NoError(t, err)
+	assert.Equal(t, stdnet.IPv4(0x0A, 0x0C, 0x71, 0x0C), ip)
+
+	ip, err = n.StdParseIP("2001:470:20::2")
+	assert.NoError(t, err)
+	assert.Equal(t, stdnet.IP{
+		0x20, 0x01, 0x04, 0x70,
+		0, 0x20, 0, 0,
+		0, 0, 0, 0,
+		0, 0, 0, 0x02,
+	}, ip)
+}
+func TestStdParseCIDR(t *testing.T) {
+	n := NetFuncs{}
+	_, err := n.StdParseCIDR("not an IP")
+	assert.Error(t, err)
+
+	_, err = n.StdParseCIDR("1.1.1.1")
+	assert.Error(t, err)
+
+	cidr, err := n.StdParseCIDR("192.168.0.2/28")
+	assert.NoError(t, err)
+	assert.Equal(t, "192.168.0.0/28", cidr.String())
+}
+
+func TestCidrHost(t *testing.T) {
+	n := NetFuncs{}
+	_, network, _ := stdnet.ParseCIDR("10.12.127.0/20")
+
+	ip, err := n.CidrHost(16, network)
+	assert.NoError(t, err)
+	assert.Equal(t, "10.12.112.16", ip.String())
+
+	ip, err = n.CidrHost(268, network)
+	assert.NoError(t, err)
+	assert.Equal(t, "10.12.113.12", ip.String())
+
+	_, network, _ = stdnet.ParseCIDR("fd00:fd12:3456:7890:00a2::/72")
+	ip, err = n.CidrHost(34, network)
+	assert.NoError(t, err)
+	assert.Equal(t, "fd00:fd12:3456:7890::22", ip.String())
+}
+
+func TestCidrNetmask(t *testing.T) {
+	n := NetFuncs{}
+	_, network, _ := stdnet.ParseCIDR("10.0.0.0/12")
+
+	ip, err := n.CidrNetmask(network)
+	assert.NoError(t, err)
+	assert.Equal(t, "255.240.0.0", ip.String())
+}
+
+func TestCidrSubnets(t *testing.T) {
+	n := NetFuncs{}
+	_, network, _ := stdnet.ParseCIDR("10.0.0.0/16")
+
+	subnets, err := n.CidrSubnets(-1, network)
+	assert.Nil(t, subnets)
+	assert.Error(t, err)
+
+	subnets, err = n.CidrSubnets(2, network)
+	assert.NoError(t, err)
+	assert.Len(t, subnets, 4)
+	assert.Equal(t, "10.0.0.0/18", subnets[0].String())
+	assert.Equal(t, "10.0.64.0/18", subnets[1].String())
+	assert.Equal(t, "10.0.128.0/18", subnets[2].String())
+	assert.Equal(t, "10.0.192.0/18", subnets[3].String())
+}
+
+func TestCidrSubnetSizes(t *testing.T) {
+	n := NetFuncs{}
+	_, network, _ := stdnet.ParseCIDR("10.1.0.0/16")
+
+	subnets, err := n.CidrSubnetSizes(network)
+	assert.Nil(t, subnets)
+	assert.Error(t, err)
+
+	subnets, err = n.CidrSubnetSizes(32, network)
+	assert.Nil(t, subnets)
+	assert.Error(t, err)
+
+	subnets, err = n.CidrSubnetSizes(-1, network)
+	assert.Nil(t, subnets)
+	assert.Error(t, err)
+
+	subnets, err = n.CidrSubnetSizes(4, 4, 8, 4, network)
+	assert.NoError(t, err)
+	assert.Len(t, subnets, 4)
+	assert.Equal(t, "10.1.0.0/20", subnets[0].String())
+	assert.Equal(t, "10.1.16.0/20", subnets[1].String())
+	assert.Equal(t, "10.1.32.0/24", subnets[2].String())
+	assert.Equal(t, "10.1.48.0/20", subnets[3].String())
+}