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

5 files changed, 278 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())
+}
diff --git a/go.mod b/go.mod
index ff133e28..2ef6dc60 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/Masterminds/goutils v1.1.1
 	github.com/Shopify/ejson v1.3.3
+	github.com/apparentlymart/go-cidr v1.1.0
 	github.com/aws/aws-sdk-go v1.44.4
 	github.com/docker/libkv v0.2.2-0.20180912205406-458977154600
 	github.com/fullsailor/pkcs7 v0.0.0-20190404230743-d7302db945fa
diff --git a/go.sum b/go.sum
index 7f3ddd32..cf6df941 100644
--- a/go.sum
+++ b/go.sum
@@ -134,6 +134,8 @@ github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRF
 github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239 h1:kFOfPq6dUM1hTo4JG6LR5AXSUEsOjtdm0kw0FtQtMJA=
 github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239/go.mod h1:2FmKhYUyUczH0OGQWaF5ceTx0UBShxjsH6f8oGKYe2c=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
+github.com/apparentlymart/go-cidr v1.1.0 h1:2mAhrMoF+nhXqxTzSZMUzDHkLjmIHC+Zzn4tdgBZjnU=
+github.com/apparentlymart/go-cidr v1.1.0/go.mod h1:EBcsNrHc3zQeuaeCeCtQruQm+n9/YjEn/vI25Lg7Gwc=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
 github.com/armon/go-metrics v0.3.9/go.mod h1:4O98XIr/9W0sxpJ8UaYkvjk10Iff7SnFrb4QAOwNTFc=
diff --git a/internal/tests/integration/net_test.go b/internal/tests/integration/net_test.go
index 526461f3..f56f0318 100644
--- a/internal/tests/integration/net_test.go
+++ b/internal/tests/integration/net_test.go
@@ -7,3 +7,40 @@ import (
 func TestNet_LookupIP(t *testing.T) {
 	inOutTest(t, `{{ net.LookupIP "localhost" }}`, "127.0.0.1")
 }
+
+func TestNet_ParseCIDR(t *testing.T) {
+	inOutTest(t, `{{ net.StdParseCIDR "10.12.127.0/20" }}`, "10.12.112.0/20")
+}
+
+func TestNet_CidrHost(t *testing.T) {
+	inOutTest(t, `{{ net.StdParseCIDR "10.12.127.0/20" | net.CidrHost 16 }}`, "10.12.112.16")
+	inOutTest(t, `{{ "10.12.127.0/20" | net.CidrHost 16 }}`, "10.12.112.16")
+	inOutTest(t, `{{ net.CidrHost 268 "10.12.127.0/20" }}`, "10.12.113.12")
+	inOutTest(t, `{{ net.CidrHost 34 "fd00:fd12:3456:7890:00a2::/72" }}`, "fd00:fd12:3456:7890::22")
+}
+
+func TestNet_CidrNetmask(t *testing.T) {
+	inOutTest(t, `{{ "10.12.127.0/20" | net.CidrNetmask }}`, "255.255.240.0")
+	inOutTest(t, `{{ net.CidrNetmask "10.0.0.0/12" }}`, "255.240.0.0")
+}
+
+func TestNet_CidrSubnets(t *testing.T) {
+	inOutTest(t, `{{ index ("10.0.0.0/16" | net.CidrSubnets 2) 1 }}`, "10.0.64.0/18")
+	inOutTest(t, `{{ range net.CidrSubnets 2 "10.0.0.0/16" }}
+{{ . }}{{ end }}`, `
+10.0.0.0/18
+10.0.64.0/18
+10.0.128.0/18
+10.0.192.0/18`)
+}
+
+func TestNet_CidrSubnetSizes(t *testing.T) {
+	inOutTest(t, `{{ index ("10.0.0.0/16" | net.CidrSubnetSizes 1) 0 }}`, "10.0.0.0/17")
+	inOutTest(t, `{{ index ("10.1.0.0/16" | net.CidrSubnetSizes 4 4 8 4) 1 }}`, "10.1.16.0/20")
+	inOutTest(t, `{{ range net.CidrSubnetSizes 16 16 16 32 "fd00:fd12:3456:7890::/56" }}
+{{ . }}{{ end }}`, `
+fd00:fd12:3456:7800::/72
+fd00:fd12:3456:7800:100::/72
+fd00:fd12:3456:7800:200::/72
+fd00:fd12:3456:7800:300::/88`)
+}