internal/agent/dhcp/dhcp_test.go

Ref: main v0.0.5   Size: 4.8 KiB   History

package dhcp

import (
	"net"
	"testing"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func testServer() *Server {
	return NewServer("eitri0", "10.77.1.0/24",
		net.IPv4(10, 77, 1, 1),
		net.CIDRMask(24, 32),
		[]net.IP{net.IPv4(1, 1, 1, 1)},
		12*time.Hour)
}

func newTestServer() *Server {
	return NewServer("eitri0", "10.77.1.0/24",
		net.IPv4(10, 77, 1, 1), net.CIDRMask(24, 32), nil, time.Hour)
}

func TestReserveAllocatesDistinctStickyAddresses(t *testing.T) {
	s := newTestServer()
	macA, _ := net.ParseMAC("52:54:00:aa:aa:aa")
	macB, _ := net.ParseMAC("52:54:00:bb:bb:bb")

	ipA, err := s.Reserve(macA)
	require.NoError(t, err)
	assert.Equal(t, "10.77.1.2", ipA.String(), "first free skips network(.0) + gateway(.1)")

	ipB, err := s.Reserve(macB)
	require.NoError(t, err)
	assert.Equal(t, "10.77.1.3", ipB.String(), "second VM gets a distinct address")

	ipA2, err := s.Reserve(macA)
	require.NoError(t, err)
	assert.Equal(t, ipA.String(), ipA2.String(), "re-reserving a MAC is sticky")
}

func TestReserveSkipsPreexistingReservation(t *testing.T) {
	s := newTestServer()
	pinned, _ := net.ParseMAC("52:54:00:cc:cc:cc")
	s.SetReservation(pinned, net.ParseIP("10.77.1.2"))

	fresh, _ := net.ParseMAC("52:54:00:dd:dd:dd")
	ip, err := s.Reserve(fresh)
	require.NoError(t, err)
	assert.Equal(t, "10.77.1.3", ip.String(), "an externally-set reservation is treated as used")
}

// TestReserveHandsBackAPreloadedReservation pins the handoff the startup replay
// depends on: an address preloaded into the table is what Reserve returns for
// its MAC, so a surviving guest keeps its address across an agent restart
// instead of being renumbered on its next boot.
func TestReserveHandsBackAPreloadedReservation(t *testing.T) {
	s := newTestServer()
	survivor, _ := net.ParseMAC("52:54:00:cc:cc:cc")
	s.SetReservation(survivor, net.ParseIP("10.77.1.55"))

	ip, err := s.Reserve(survivor)
	require.NoError(t, err)
	assert.Equal(t, "10.77.1.55", ip.String(), "a preloaded reservation is never renumbered")
}

func mustMAC(t *testing.T, s string) net.HardwareAddr {
	t.Helper()
	m, err := net.ParseMAC(s)
	if err != nil {
		t.Fatalf("parse mac %q: %v", s, err)
	}
	return m
}

func TestReservationSetLookupRemove(t *testing.T) {
	s := testServer()
	mac := mustMAC(t, "52:54:00:ab:cd:ef")

	if _, ok := s.lookup(mac); ok {
		t.Fatal("unknown MAC should not resolve")
	}
	s.SetReservation(mac, net.IPv4(10, 77, 1, 7))
	got, ok := s.lookup(mac)
	if !ok || !got.Equal(net.IPv4(10, 77, 1, 7)) {
		t.Fatalf("lookup after set = %v, %v; want 10.77.1.7, true", got, ok)
	}
	s.RemoveReservation(mac)
	if _, ok := s.lookup(mac); ok {
		t.Fatal("lookup after remove should be false")
	}
}

func TestReservationLookupIsCaseInsensitiveOnMAC(t *testing.T) {
	s := testServer()
	s.SetReservation(mustMAC(t, "52:54:00:AB:CD:EF"), net.IPv4(10, 77, 1, 9))
	if _, ok := s.lookup(mustMAC(t, "52:54:00:ab:cd:ef")); !ok {
		t.Fatal("reservation lookup must not depend on MAC hex case")
	}
}

func request(t *testing.T, mac net.HardwareAddr, mt dhcpv4.MessageType) *dhcpv4.DHCPv4 {
	t.Helper()
	m, err := dhcpv4.New(dhcpv4.WithHwAddr(mac), dhcpv4.WithMessageType(mt))
	if err != nil {
		t.Fatalf("build request: %v", err)
	}
	return m
}

func TestBuildReplyUnknownMACGetsNothing(t *testing.T) {
	s := testServer()
	resp, err := s.buildReply(request(t, mustMAC(t, "52:54:00:11:22:33"), dhcpv4.MessageTypeDiscover))
	if err != nil {
		t.Fatalf("buildReply err: %v", err)
	}
	if resp != nil {
		t.Fatalf("unknown MAC must get no reply, got %v", resp)
	}
}

func TestBuildReplyDiscoverOffersReservedIP(t *testing.T) {
	s := testServer()
	mac := mustMAC(t, "52:54:00:ab:cd:ef")
	s.SetReservation(mac, net.IPv4(10, 77, 1, 7))

	resp, err := s.buildReply(request(t, mac, dhcpv4.MessageTypeDiscover))
	if err != nil {
		t.Fatalf("buildReply err: %v", err)
	}
	if resp == nil {
		t.Fatal("reserved MAC must get a reply")
	}
	if resp.MessageType() != dhcpv4.MessageTypeOffer {
		t.Fatalf("Discover must yield Offer, got %v", resp.MessageType())
	}
	if !resp.YourIPAddr.Equal(net.IPv4(10, 77, 1, 7)) {
		t.Fatalf("YourIPAddr = %v; want 10.77.1.7", resp.YourIPAddr)
	}
	if !resp.Router()[0].Equal(net.IPv4(10, 77, 1, 1)) {
		t.Fatalf("router = %v; want gateway 10.77.1.1", resp.Router())
	}
	if len(resp.DNS()) == 0 || !resp.DNS()[0].Equal(net.IPv4(1, 1, 1, 1)) {
		t.Fatalf("DNS = %v; want [1.1.1.1]", resp.DNS())
	}
}

func TestBuildReplyRequestAcks(t *testing.T) {
	s := testServer()
	mac := mustMAC(t, "52:54:00:ab:cd:ef")
	s.SetReservation(mac, net.IPv4(10, 77, 1, 7))
	resp, err := s.buildReply(request(t, mac, dhcpv4.MessageTypeRequest))
	if err != nil {
		t.Fatalf("buildReply err: %v", err)
	}
	if resp == nil || resp.MessageType() != dhcpv4.MessageTypeAck {
		t.Fatalf("Request must yield Ack, got %v", resp)
	}
}