Add some basic dhcp tests

internal/data/dhcp.go

@@ -2,25 +2,40 @@ package data
 
 import (
 	"context"
+	"encoding/binary"
 	"errors"
 	"math"
+	"math/big"
 	"math/rand"
 	"net"
 
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/client/v3/clientv3util"
 )
 
-// https://gist.github.com/udhos/b468fbfd376aa0b655b6b0c539a88c03
 func incrementIP(ip net.IP, inc uint) net.IP {
-	i := ip.To4()
-	v := uint(i[0])<<24 + uint(i[1])<<16 + uint(i[2])<<8 + uint(i[3])
-	v += inc
-	v3 := byte(v & 0xFF)
-	v2 := byte((v >> 8) & 0xFF)
-	v1 := byte((v >> 16) & 0xFF)
-	v0 := byte((v >> 24) & 0xFF)
-	return net.IPv4(v0, v1, v2, v3)
+
+	if ip.To4() != nil {
+		r := binary.BigEndian.Uint32(ip.To4()) + uint32(inc)
+
+		newIp := make([]byte, 4)
+		binary.BigEndian.PutUint32(newIp, r)
+
+		return net.IP(newIp)
+	}
+
+	ip = ip.To16()
+
+	asBigInt := big.NewInt(0).SetBytes(ip)
+	asBigInt.Add(asBigInt, big.NewInt(int64(inc)))
+
+	result := make([]byte, 16)
+	bigIntBytes := asBigInt.Bytes()
+
+	copy(result[16-min(len(bigIntBytes), 16):], bigIntBytes)
+
+	return net.IP(result)
+
 }
 
 func getNextIP(subnet string) (string, error) {
@@ -30,8 +45,13 @@ func getNextIP(subnet string) (string, error) {
 		return "", err
 	}
 
+	max := 32
+	if serverIP.To16() != nil {
+		max = 128
+	}
+
 	used, _ := cidr.Mask.Size()
-	maxNumberOfAddresses := int(math.Pow(2, float64(32-used))) - 2 // Do not allocate largest address or 0
+	maxNumberOfAddresses := int(math.Pow(2, float64(max-used))) - 2 // Do not allocate largest address or 0
 	if maxNumberOfAddresses < 1 {
 		return "", errors.New("subnet is too small to contain a new device")
 	}

internal/data/dhcp_test.go

@@ -0,0 +1,133 @@
+package data
+
+import (
+	"log"
+	"net"
+	"testing"
+)
+
+func TestIncrementIP(t *testing.T) {
+	tests := []struct {
+		name     string
+		ip       string
+		inc      uint
+		expected string
+	}{
+		// IPv4 tests
+		{
+			name:     "IPv4 increment by 1",
+			ip:       "192.168.1.1",
+			inc:      1,
+			expected: "192.168.1.2",
+		},
+		{
+			name:     "IPv4 increment by 255",
+			ip:       "192.168.1.1",
+			inc:      255,
+			expected: "192.168.2.0",
+		},
+		{
+			name:     "IPv4 crosses multiple octets",
+			ip:       "192.168.255.255",
+			inc:      1,
+			expected: "192.169.0.0",
+		},
+		// IPv6 tests
+		{
+			name:     "IPv6 increment by 1",
+			ip:       "2001:db8::1",
+			inc:      1,
+			expected: "2001:db8::2",
+		},
+		{
+			name:     "IPv6 increment by 255",
+			ip:       "2001:db8::1",
+			inc:      255,
+			expected: "2001:db8::100",
+		},
+		{
+			name:     "IPv6 increment across boundary",
+			ip:       "2001:db8::ffff",
+			inc:      1,
+			expected: "2001:db8::1:0",
+		},
+		{
+			name:     "IPv6 increment zero address",
+			ip:       "::",
+			inc:      1,
+			expected: "::1",
+		},
+		{
+			name:     "IPv6 increment large number",
+			ip:       "2001:db8::",
+			inc:      65535,
+			expected: "2001:db8::ffff",
+		},
+		{
+			name:     "IPv6 with zeros in middle",
+			ip:       "2001:db8:0:0:0:0:0:1",
+			inc:      1,
+			expected: "2001:db8::2",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ip := net.ParseIP(tt.ip)
+			if ip == nil {
+				t.Fatalf("failed to parse IP: %s", tt.ip)
+			}
+
+			result := incrementIP(ip, tt.inc)
+			if result.String() != tt.expected {
+
+				log.Println(len(ip), len(result))
+
+				t.Errorf("incrementIP(%s, %d) = %s; want %s",
+					tt.ip, tt.inc, result.String(), tt.expected)
+			}
+		})
+	}
+}
+
+func TestIncrementIPOverflow(t *testing.T) {
+	tests := []struct {
+		name string
+		ip   string
+		inc  uint
+	}{
+		{
+			name: "IPv4 overflow max IP",
+			ip:   "255.255.255.255",
+			inc:  1,
+		},
+		{
+			name: "IPv6 overflow max IP",
+			ip:   "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff",
+			inc:  1,
+		},
+		{
+			name: "IPv6 large overflow",
+			ip:   "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ff00",
+			inc:  256,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ip := net.ParseIP(tt.ip)
+			if ip == nil {
+				t.Fatalf("failed to parse IP: %s", tt.ip)
+			}
+
+			defer func() {
+				if r := recover(); r != nil {
+					t.Errorf("incrementIP(%s, %d) paniced",
+						tt.ip, tt.inc)
+				}
+			}()
+
+			incrementIP(ip, tt.inc)
+		})
+	}
+}