flock/pkg/agent/ipam_test.go

package agent

import (
	"net"
	"testing"

	"code.fritzlab.net/fritzlab/flock/pkg/embed"
)

// fakeRand is a deterministic randSource for tests.
type fakeRand struct {
	nibbles []byte // queue of NibbleN results
	iids    [][]byte
	picks   []int
	n, i, p int
}

func (f *fakeRand) NibbleN() byte {
	v := f.nibbles[f.n%len(f.nibbles)]
	f.n++
	return v & 0x0F
}
func (f *fakeRand) FillIID(dst []byte) {
	src := f.iids[f.i%len(f.iids)]
	f.i++
	copy(dst, src)
}
func (f *fakeRand) PickIndex(n int) int {
	if len(f.picks) == 0 {
		return 0
	}
	v := f.picks[f.p%len(f.picks)]
	f.p++
	if v >= n {
		return 0
	}
	return v
}

func mustNet(t *testing.T, s string) *net.IPNet {
	t.Helper()
	_, n, err := net.ParseCIDR(s)
	if err != nil {
		t.Fatal(err)
	}
	return n
}

func TestIntersectCIDR(t *testing.T) {
	v6a := mustNet(t, "2602:817:3000:f001::/64")
	v6super := mustNet(t, "2602:817:3000::/48")
	v6sub := mustNet(t, "2602:817:3000:f001::/96")
	v6other := mustNet(t, "2602:817:3000:f002::/64")
	v4 := mustNet(t, "10.0.0.0/24")

	cases := []struct {
		name string
		ann  *net.IPNet
		node *net.IPNet
		want *net.IPNet
	}{
		{"equal → node", v6a, v6a, v6a},
		{"supernet of node → node", v6super, v6a, v6a},
		{"subnet of node → ann", v6sub, v6a, v6sub},
		{"disjoint same family → nil", v6other, v6a, nil},
		{"different family → nil", v4, v6a, nil},
	}
	for _, c := range cases {
		got := intersectCIDR(c.ann, c.node)
		if (got == nil) != (c.want == nil) {
			t.Errorf("%s: got=%v want=%v", c.name, got, c.want)
			continue
		}
		if got != nil && !cidrEqual(got, c.want) {
			t.Errorf("%s: got=%s want=%s", c.name, got, c.want)
		}
	}
}

func TestResolveEffective(t *testing.T) {
	nodes := []*net.IPNet{
		mustNet(t, "2602:817:3000:f001::/64"),
		mustNet(t, "2602:817:3000:f002::/64"),
	}

	// Empty annotation → all node CIDRs.
	if got, err := resolveEffective(nil, nodes); err != nil || len(got) != 2 {
		t.Fatalf("empty ann: got=%v err=%v", got, err)
	}

	// Subnet of one node CIDR → that subnet only.
	ann := []*net.IPNet{mustNet(t, "2602:817:3000:f001::/96")}
	got, err := resolveEffective(ann, nodes)
	if err != nil || len(got) != 1 || got[0].String() != "2602:817:3000:f001::/96" {
		t.Fatalf("subnet ann: got=%v err=%v", got, err)
	}

	// Supernet matches both node CIDRs.
	ann = []*net.IPNet{mustNet(t, "2602:817:3000::/48")}
	got, err = resolveEffective(ann, nodes)
	if err != nil || len(got) != 2 {
		t.Fatalf("supernet ann: got=%v err=%v", got, err)
	}

	// First-match-wins: the first ann CIDR that intersects anything wins.
	ann = []*net.IPNet{
		mustNet(t, "2602:817:3000:ff::/64"), // no intersection
		mustNet(t, "2602:817:3000:f001::/64"),
	}
	got, err = resolveEffective(ann, nodes)
	if err != nil || len(got) != 1 || got[0].String() != "2602:817:3000:f001::/64" {
		t.Fatalf("first-match: got=%v err=%v", got, err)
	}

	// No overlap → error.
	ann = []*net.IPNet{mustNet(t, "2602:817:3000:ff::/64")}
	if _, err := resolveEffective(ann, nodes); err == nil {
		t.Fatalf("expected error for disjoint ann")
	}
}

func TestIPAM_AllocV6_Random(t *testing.T) {
	i, err := NewIPAM([]string{"2602:817:3000:f001::/64"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	i.randSrc = &fakeRand{
		iids: [][]byte{
			// 16 bytes; only the low 8 bytes are used for a /64 host portion.
			{0, 0, 0, 0, 0, 0, 0, 0, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x11, 0x22},
		},
	}
	res, err := i.Allocate(AllocRequest{
		ContainerID: "c1", Namespace: "ns", Pod: "p", WantV6: true,
	})
	if err != nil {
		t.Fatalf("Allocate: %v", err)
	}
	if res.IP6.String() != "2602:817:3000:f001:aabb:ccdd:eeff:1122" {
		t.Fatalf("IP6 = %s, want 2602:817:3000:f001:aabb:ccdd:eeff:1122", res.IP6)
	}
}

func TestIPAM_AllocV6_WithEmbed(t *testing.T) {
	i, err := NewIPAM([]string{"2602:817:3000:f001::/64"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	i.randSrc = &fakeRand{nibbles: []byte{0xe}}
	res, err := i.Allocate(AllocRequest{
		ContainerID: "c1", Namespace: "mail", Pod: "stalwart-0", App: "stalwart", WantV6: true,
		IPAlgo: []embed.Field{embed.FieldNamespace, embed.FieldApp, embed.FieldImage},
	})
	if err != nil {
		t.Fatalf("Allocate: %v", err)
	}
	// Last nibble from fakeRand NibbleN → 0xe.
	if got := res.IP6[15] & 0x0F; got != 0xe {
		t.Fatalf("last nibble = %x, want e", got)
	}
	mustNet(t, "2602:817:3000:f001::/64").Contains(res.IP6)
}

func TestIPAM_AllocV6_CollisionRetry(t *testing.T) {
	i, err := NewIPAM([]string{"2602:817:3000:f001::/64"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	// Mark one specific address in-use, then feed the allocator that same
	// address first and a new one on retry.
	first := net.ParseIP("2602:817:3000:f001::1")
	second := net.ParseIP("2602:817:3000:f001::2")
	i.MarkInUse(first)
	i.randSrc = &fakeRand{
		iids: [][]byte{
			append(make([]byte, 8), 0, 0, 0, 0, 0, 0, 0, 0x01), // -> ...:1 (collides)
			append(make([]byte, 8), 0, 0, 0, 0, 0, 0, 0, 0x02), // -> ...:2 (ok)
		},
	}
	res, err := i.Allocate(AllocRequest{
		ContainerID: "c1", Namespace: "ns", Pod: "p", WantV6: true,
	})
	if err != nil {
		t.Fatalf("Allocate: %v", err)
	}
	if !res.IP6.Equal(second) {
		t.Fatalf("IP6 = %s, want %s (collision retry)", res.IP6, second)
	}
}

func TestIPAM_AllocV4_SkipsNetworkAndGateway(t *testing.T) {
	i, err := NewIPAM(nil, []string{"172.25.210.0/24"})
	if err != nil {
		t.Fatal(err)
	}
	i.randSrc = &fakeRand{}
	res, err := i.Allocate(AllocRequest{
		ContainerID: "c1", Namespace: "ns", Pod: "p", WantV4: true,
	})
	if err != nil {
		t.Fatalf("Allocate: %v", err)
	}
	if res.IP4.String() != "172.25.210.2" {
		t.Fatalf("IP4 = %s, want 172.25.210.2 (skip .0 network + .1 gateway)", res.IP4)
	}
}

func TestIPAM_AllocV4_Sequential(t *testing.T) {
	i, err := NewIPAM(nil, []string{"172.25.210.0/29"})
	if err != nil {
		t.Fatal(err)
	}
	i.randSrc = &fakeRand{}
	want := []string{"172.25.210.2", "172.25.210.3", "172.25.210.4", "172.25.210.5", "172.25.210.6"}
	for _, w := range want {
		res, err := i.Allocate(AllocRequest{ContainerID: w, WantV4: true})
		if err != nil {
			t.Fatalf("Allocate: %v", err)
		}
		if res.IP4.String() != w {
			t.Fatalf("got %s, want %s", res.IP4, w)
		}
	}
	// Next allocation should fail — /29 exhausted (.0 network, .1 gateway, .7 broadcast).
	if _, err := i.Allocate(AllocRequest{ContainerID: "extra", WantV4: true}); err == nil {
		t.Fatalf("expected exhaustion error")
	}
}

func TestIPAM_DualStack(t *testing.T) {
	i, err := NewIPAM(
		[]string{"2602:817:3000:f001::/64"},
		[]string{"172.25.210.0/24"},
	)
	if err != nil {
		t.Fatal(err)
	}
	i.randSrc = &fakeRand{
		iids: [][]byte{
			{0, 0, 0, 0, 0, 0, 0, 0, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x11, 0x22},
		},
	}
	res, err := i.Allocate(AllocRequest{
		ContainerID: "c1", WantV6: true, WantV4: true,
	})
	if err != nil {
		t.Fatalf("Allocate: %v", err)
	}
	if res.IP6 == nil || res.IP4 == nil {
		t.Fatalf("dual-stack result missing IPs: %+v", res)
	}
	if res.IP4.String() != "172.25.210.2" {
		t.Fatalf("IP4 = %s", res.IP4)
	}
}

func TestIPAM_RejectsNoFamily(t *testing.T) {
	i, _ := NewIPAM([]string{"2602:817:3000:f001::/64"}, nil)
	if _, err := i.Allocate(AllocRequest{ContainerID: "c"}); err == nil {
		t.Fatalf("expected error for no v6/v4")
	}
}

func TestIPAM_Release(t *testing.T) {
	i, _ := NewIPAM([]string{"2602:817:3000:f001::/64"}, nil)
	i.randSrc = &fakeRand{
		iids: [][]byte{
			{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},
			{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},
		},
	}
	res1, err := i.Allocate(AllocRequest{ContainerID: "c1", WantV6: true})
	if err != nil {
		t.Fatal(err)
	}
	i.Release(res1.IP6)
	// Next Allocate should be free to pick the same address again.
	res2, err := i.Allocate(AllocRequest{ContainerID: "c2", WantV6: true})
	if err != nil {
		t.Fatal(err)
	}
	if !res2.IP6.Equal(res1.IP6) {
		t.Fatalf("expected release to free %s; got %s", res1.IP6, res2.IP6)
	}
}

func TestIPAM_RejectsMismatchedFamily(t *testing.T) {
	if _, err := NewIPAM([]string{"10.0.0.0/24"}, nil); err == nil {
		t.Fatalf("expected v4 in cidr6 to fail")
	}
	if _, err := NewIPAM(nil, []string{"2602:817:3000::/64"}); err == nil {
		t.Fatalf("expected v6 in cidr4 to fail")
	}
}