a73x

internal/agent/exposeproxy/udp_test.go

Ref:   Size: 16.0 KiB   History

package exposeproxy

import (
	"bytes"
	"net"
	"testing"
	"time"

	"github.com/a73x/eitri/internal/pb"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// fakeUDPGuest is a bound UDP socket standing in for a service inside a guest.
// It echoes every datagram back to whoever sent it, prefixed, so a test can
// prove a datagram went both ways through the proxy and came back whole.
type fakeUDPGuest struct {
	pc   *net.UDPConn
	addr string
	port uint32
}

func newFakeUDPGuest(t *testing.T) *fakeUDPGuest {
	t.Helper()
	pc, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	require.NoError(t, err)
	t.Cleanup(func() { pc.Close() })
	g := &fakeUDPGuest{pc: pc, addr: "127.0.0.1", port: uint32(pc.LocalAddr().(*net.UDPAddr).Port)}
	go func() {
		buf := make([]byte, 64*1024)
		for {
			n, from, err := pc.ReadFromUDP(buf)
			if err != nil {
				return
			}
			pc.WriteToUDP(append([]byte("echo:"), buf[:n]...), from) //nolint:errcheck
		}
	}()
	return g
}

// boundUDPPort reads the address the manager's packet socket for id bound.
func boundUDPPort(t *testing.T, m *Manager, id string) string {
	t.Helper()
	m.mu.Lock()
	defer m.mu.Unlock()
	ex, ok := m.live[id]
	require.True(t, ok, "no exposure %q", id)
	require.NotNil(t, ex.pc, "exposure %q has no packet socket", id)
	return ex.pc.LocalAddr().String()
}

// udpClient is one caller of a published UDP port, held open across datagrams
// so the proxy sees the same source address each time — which is what makes it
// one session rather than several.
func udpClient(t *testing.T, addr string) *net.UDPConn {
	t.Helper()
	ua, err := net.ResolveUDPAddr("udp4", addr)
	require.NoError(t, err)
	c, err := net.DialUDP("udp4", nil, ua)
	require.NoError(t, err)
	t.Cleanup(func() { c.Close() })
	require.NoError(t, c.SetDeadline(time.Now().Add(5*time.Second)))
	return c
}

// say sends one datagram and returns the one that comes back.
func say(t *testing.T, c *net.UDPConn, msg string) string {
	t.Helper()
	_, err := c.Write([]byte(msg))
	require.NoError(t, err)
	buf := make([]byte, 64*1024)
	n, err := c.Read(buf)
	require.NoError(t, err)
	return string(buf[:n])
}

// sessions is how many the exposure is holding right now.
func sessions(t *testing.T, m *Manager, id string) int {
	t.Helper()
	m.mu.Lock()
	ex, ok := m.live[id]
	m.mu.Unlock()
	require.True(t, ok, "no exposure %q", id)
	ex.smu.Lock()
	defer ex.smu.Unlock()
	return len(ex.sessions)
}

// waitSessions waits for the session count to reach want. A session is filed
// synchronously with the datagram that starts it, but it is torn down by a
// goroutine finishing on its own time.
func waitSessions(t *testing.T, m *Manager, id string, want int) {
	t.Helper()
	var got int
	for i := 0; i < 200; i++ {
		if got = sessions(t, m, id); got == want {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("exposure %q holds %d sessions, want %d", id, got, want)
}

// oneSession returns the exposure's single session, for tests that reach past
// the wire to what the proxy is holding.
func oneSession(t *testing.T, m *Manager, id string) *udpSession {
	t.Helper()
	m.mu.Lock()
	ex, ok := m.live[id]
	m.mu.Unlock()
	require.True(t, ok, "no exposure %q", id)
	ex.smu.Lock()
	defer ex.smu.Unlock()
	require.Len(t, ex.sessions, 1, "want exactly one session")
	for _, s := range ex.sessions {
		return s
	}
	return nil
}

func TestUDPConvergeBindsAndForwardsBothWays(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})

	got := m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	require.Len(t, got, 1)
	assert.Equal(t, "active", got[0].GetState(), "active means the packet socket is bound")

	c := udpClient(t, boundUDPPort(t, m, "e1"))
	assert.Equal(t, "echo:hello", say(t, c, "hello"))
	assert.Equal(t, 1, sessions(t, m, "e1"), "one client address is one session")
}

func TestUDPKeepsDatagramBoundaries(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	c := udpClient(t, boundUDPPort(t, m, "e1"))

	// Two datagrams in, two out, in order and whole — nothing here coalesces a
	// stream out of them.
	assert.Equal(t, "echo:one", say(t, c, "one"))
	assert.Equal(t, "echo:two", say(t, c, "two"))

	// And a big one survives intact: the read buffers are sized past the
	// largest datagram either side can send.
	big := string(bytes.Repeat([]byte("x"), 8000))
	assert.Equal(t, "echo:"+big, say(t, c, big))
}

func TestUDPSessionIsPerClientAddress(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	addr := boundUDPPort(t, m, "e1")

	first, second := udpClient(t, addr), udpClient(t, addr)
	assert.Equal(t, "echo:a", say(t, first, "a"))
	assert.Equal(t, "echo:b", say(t, second, "b"))
	assert.Equal(t, 2, sessions(t, m, "e1"), "two callers are two conversations")

	// The same caller again is the same session, not a third.
	assert.Equal(t, "echo:again", say(t, first, "again"))
	assert.Equal(t, 2, sessions(t, m, "e1"))
}

func TestUDPSessionIsPromotedOnlyByAGuestReply(t *testing.T) {
	// A guest that takes datagrams and never answers: the session stays on the
	// short window, however much the client says.
	silent, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	require.NoError(t, err)
	t.Cleanup(func() { silent.Close() })
	quietPort := uint32(silent.LocalAddr().(*net.UDPAddr).Port)

	m := newTestManager(t, map[string]string{"vm1": "127.0.0.1"})
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", quietPort, 0)})
	c := udpClient(t, boundUDPPort(t, m, "e1"))
	_, err = c.Write([]byte("anyone there"))
	require.NoError(t, err)
	waitSessions(t, m, "e1", 1)

	s := oneSession(t, m, "e1")
	assert.False(t, s.replied.Load(), "nothing has come back; the session is still a guess")
	assert.WithinDuration(t, time.Now().Add(m.unrepliedIdle), s.expiry(), time.Second)

	// Now a guest that answers. The reply is the promotion.
	g := newFakeUDPGuest(t)
	m.Converge([]*pb.ExposureSpec{desiredUDP("e2", "vm1", g.port, 0)})
	talking := udpClient(t, boundUDPPort(t, m, "e2"))
	require.Equal(t, "echo:hi", say(t, talking, "hi"))

	promoted := oneSession(t, m, "e2")
	assert.True(t, promoted.replied.Load())
	assert.WithinDuration(t, time.Now().Add(m.repliedIdle), promoted.expiry(), time.Second)
}

func TestUDPSessionExpiresOnBothWindows(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	// Windows a test can outlast. Everything about them is what the shipped
	// constants do, at a scale that fits in a test.
	m.unrepliedIdle, m.repliedIdle = 250*time.Millisecond, 10*time.Second

	// Never answered: gone at the short window.
	silent, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	require.NoError(t, err)
	t.Cleanup(func() { silent.Close() })
	m.Converge([]*pb.ExposureSpec{
		desiredUDP("quiet", "vm1", uint32(silent.LocalAddr().(*net.UDPAddr).Port), 0),
		desiredUDP("live", "vm1", g.port, 0),
	})

	unanswered := udpClient(t, boundUDPPort(t, m, "quiet"))
	_, err = unanswered.Write([]byte("hello?"))
	require.NoError(t, err)
	waitSessions(t, m, "quiet", 1)
	waitSessions(t, m, "quiet", 0)

	// Answered once: the same stretch of quiet does not end it, because the
	// promotion bought it the longer window.
	c := udpClient(t, boundUDPPort(t, m, "live"))
	require.Equal(t, "echo:hi", say(t, c, "hi"))
	waitSessions(t, m, "live", 1)
	time.Sleep(600 * time.Millisecond)
	assert.Equal(t, 1, sessions(t, m, "live"), "a conversation that has answered outlives the short window")
	assert.Equal(t, "echo:still here", say(t, c, "still here"), "and it is the same session, still carrying traffic")
}

func TestUDPRefusesNewSessionsAtItsCapWithoutEvicting(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	m.maxSessions = 1
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	addr := boundUDPPort(t, m, "e1")

	held := udpClient(t, addr)
	require.Equal(t, "echo:mine", say(t, held, "mine"), "the first caller has the only slot, and has been answered")

	// The table is full. The second caller's datagram is dropped where it
	// arrives: no session, and nothing comes back.
	over := udpClient(t, addr)
	_, err := over.Write([]byte("me too"))
	require.NoError(t, err)
	require.NoError(t, over.SetReadDeadline(time.Now().Add(300*time.Millisecond)))
	buf := make([]byte, 64)
	_, err = over.Read(buf)
	assert.Error(t, err, "a caller past the cap is answered by nothing")
	assert.Equal(t, 1, sessions(t, m, "e1"))

	// And the conversation that was already working is untouched — a full table
	// refuses the new, it does not sacrifice the live.
	assert.Equal(t, "echo:still mine", say(t, held, "still mine"))
}

func TestUDPConvergeEvictsASessionWhoseGuestMoved(t *testing.T) {
	first, second := newFakeUDPGuest(t), newFakeUDPGuest(t)
	where := map[string]string{"vm1": first.addr}
	m := NewManager(func(vmID string) string { return where[vmID] })
	t.Cleanup(m.StopAll)

	// Both fakes answer on 127.0.0.1, so the guest PORT is what tells them
	// apart; the pin under test is the address, so move the VM to an address
	// nothing is at and prove the session goes.
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", first.port, 0)})
	c := udpClient(t, boundUDPPort(t, m, "e1"))
	require.Equal(t, "echo:before", say(t, c, "before"))
	s := oneSession(t, m, "e1")

	where["vm1"] = "127.0.0.2"
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", first.port, 0)})
	assert.Equal(t, 0, sessions(t, m, "e1"), "a pin that no longer matches is not a session to keep")
	_, err := s.guest.Write([]byte("orphan"))
	assert.Error(t, err, "the evicted session's socket is closed, not leaked")

	// A guest the host has lost track of entirely is no different: there is
	// nowhere to send, so there is no session.
	where["vm1"] = second.addr
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", second.port, 0)})
	require.Equal(t, "echo:after", say(t, udpClient(t, boundUDPPort(t, m, "e1")), "after"))
	where["vm1"] = ""
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", second.port, 0)})
	assert.Equal(t, 0, sessions(t, m, "e1"), "no address is not an address to keep sending to")
}

func TestUDPSessionEndsWhenItsGuestSocketFails(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	c := udpClient(t, boundUDPPort(t, m, "e1"))
	require.Equal(t, "echo:up", say(t, c, "up"))

	// Break the socket toward the guest the way an ICMP port-unreachable or a
	// descriptor going bad would: the relay's read fails, and the session that
	// cannot reach its guest stops being one.
	s := oneSession(t, m, "e1")
	require.NoError(t, s.guest.Close())
	waitSessions(t, m, "e1", 0)

	// The port is still published; the next datagram simply starts a fresh
	// session, which is what makes the failure survivable.
	assert.Equal(t, "echo:again", say(t, c, "again"))
}

func TestUDPDropsADatagramForAGuestWithNoAddress(t *testing.T) {
	m := newTestManager(t, map[string]string{}) // the guest is still leasing
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", 8080, 0)})

	c := udpClient(t, boundUDPPort(t, m, "e1"))
	_, err := c.Write([]byte("anyone"))
	require.NoError(t, err)
	require.NoError(t, c.SetReadDeadline(time.Now().Add(300*time.Millisecond)))
	buf := make([]byte, 64)
	_, err = c.Read(buf)
	assert.Error(t, err, "nowhere to send is a drop, not a wait")
	assert.Equal(t, 0, sessions(t, m, "e1"), "a session pinned to nowhere is not a session")
}

func TestUDPConvergeReportsABindFailure(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})

	squatter, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero})
	require.NoError(t, err)
	held := uint32(squatter.LocalAddr().(*net.UDPAddr).Port)

	got := m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, held)})
	require.Len(t, got, 1)
	assert.Equal(t, "failed", got[0].GetState())
	assert.NotEmpty(t, got[0].GetReason(), "the report carries what the OS said")

	require.NoError(t, squatter.Close())
	var state string
	for i := 0; i < 20; i++ {
		got = m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, held)})
		state = got[0].GetState()
		if state == "active" {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
	assert.Equal(t, "active", state, "a bind-time refusal is a level to converge on, not an error path")
}

func TestChangingTheProtocolRebinds(t *testing.T) {
	tcpGuest := newFakeGuest(t)
	udpGuest := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": tcpGuest.addr})

	m.Converge([]*pb.ExposureSpec{desired("e1", "vm1", tcpGuest.port, 0)})
	assert.Equal(t, "echo:tcp", speak(t, boundPort(t, m, "e1"), "tcp"))

	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", udpGuest.port, 0)})
	m.mu.Lock()
	ex := m.live["e1"]
	m.mu.Unlock()
	assert.Nil(t, ex.ln, "the listener the old protocol needed is gone")
	assert.Equal(t, "echo:udp", say(t, udpClient(t, boundUDPPort(t, m, "e1")), "udp"))
}

func TestStopAllClosesPacketSocketsToo(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := NewManager(func(string) string { return g.addr })
	m.Converge([]*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)})
	addr := boundUDPPort(t, m, "e1")
	c := udpClient(t, addr)
	require.Equal(t, "echo:live", say(t, c, "live"))
	s := oneSession(t, m, "e1")

	m.StopAll()

	// The socket is unbindable-again proof: the port frees, so something else
	// can take it. The session that hung off it is closed with it.
	freed, err := net.ListenUDP("udp4", mustUDPAddr(t, addr))
	require.NoError(t, err, "the published port survived StopAll")
	freed.Close()
	_, err = s.guest.Write([]byte("orphan"))
	assert.Error(t, err, "a session outliving its socket would hold a descriptor nothing owns")
}

func mustUDPAddr(t *testing.T, addr string) *net.UDPAddr {
	t.Helper()
	ua, err := net.ResolveUDPAddr("udp4", addr)
	require.NoError(t, err)
	return ua
}

// TestUDPCountsItsSessionsAndRefusals is the counters' UDP half: sessions are
// what a published UDP port holds, and a full table turns callers away with
// nothing to show for it on the wire — no connection to refuse, no error to
// return, just a datagram that goes nowhere.
func TestUDPCountsItsSessionsAndRefusals(t *testing.T) {
	g := newFakeUDPGuest(t)
	m := newTestManager(t, map[string]string{"vm1": g.addr})
	m.maxSessions = 1
	spec := []*pb.ExposureSpec{desiredUDP("e1", "vm1", g.port, 0)}
	m.Converge(spec)
	addr := boundUDPPort(t, m, "e1")

	held := udpClient(t, addr)
	require.Equal(t, "echo:mine", say(t, held, "mine"))
	waitSessions(t, m, "e1", 1)

	got := counters(t, m, spec, "e1")
	assert.Equal(t, int64(1), got.GetActive(), "a live conversation is one active session")
	assert.Zero(t, got.GetRefused())

	over := udpClient(t, addr)
	_, err := over.Write([]byte("me too"))
	require.NoError(t, err)
	require.NoError(t, over.SetReadDeadline(time.Now().Add(300*time.Millisecond)))
	_, err = over.Read(make([]byte, 64))
	require.Error(t, err, "precondition: the table is full and the datagram went nowhere")

	got = counters(t, m, spec, "e1")
	assert.Equal(t, int64(1), got.GetRefused(), "the dropped datagram is the only trace that caller left")
	assert.Equal(t, int64(1), got.GetActive(), "and the conversation that was working is untouched")
}

// TestUDPCountsADatagramForAGuestWithNoAddress: a guest still leasing is not a
// port at its cap, and the two must not read the same in a report.
func TestUDPCountsADatagramForAGuestWithNoAddress(t *testing.T) {
	m := newTestManager(t, map[string]string{}) // vm1 has no address yet
	spec := []*pb.ExposureSpec{desiredUDP("e1", "vm1", 9999, 0)}
	m.Converge(spec)
	addr := boundUDPPort(t, m, "e1")

	c := udpClient(t, addr)
	_, err := c.Write([]byte("anyone there"))
	require.NoError(t, err)
	require.NoError(t, c.SetReadDeadline(time.Now().Add(300*time.Millisecond)))
	_, err = c.Read(make([]byte, 64))
	require.Error(t, err)

	got := counters(t, m, spec, "e1")
	assert.Equal(t, int64(1), got.GetDropped())
	assert.Zero(t, got.GetRefused())
	assert.Zero(t, got.GetActive(), "a dropped datagram starts no session")
}