nebula/udp/udp_tester.go

//go:build e2e_testing

package udp

import (
	"context"
	"io"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"sync"
	"sync/atomic"
	"syscall"

	"github.com/slackhq/nebula/config"
	"github.com/slackhq/nebula/header"
)

type Packet struct {
	To   netip.AddrPort
	From netip.AddrPort
	Data []byte
}

// Copy returns a fresh *Packet (from the freelist) with a duplicate Data buffer.
func (u *Packet) Copy() *Packet {
	n := acquirePacket()
	n.To = u.To
	n.From = u.From
	if cap(n.Data) < len(u.Data) {
		n.Data = make([]byte, len(u.Data))
	} else {
		n.Data = n.Data[:len(u.Data)]
	}
	copy(n.Data, u.Data)
	return n
}

// Release returns p to the harness packet freelist.
// Callers that pull a *Packet from Get / TxPackets must Release when done.
// Channel-backed instead of sync.Pool because sync.Pool's per-P caches drain badly under cross-goroutine Get/Put,
// and putting a []byte in a Pool escapes the slice header to heap.
func (p *Packet) Release() {
	if p == nil {
		return
	}
	p.Data = p.Data[:0]
	select {
	case packetFreelist <- p:
	default:
		// Freelist full; drop the *Packet for the GC.
	}
}

// packetFreelist retains *Packet structs (and their backing Data arrays) so steady-state allocation drops to zero.
var packetFreelist = make(chan *Packet, 64)

func acquirePacket() *Packet {
	select {
	case p := <-packetFreelist:
		return p
	default:
		return &Packet{}
	}
}

// testerAddrs is where a TesterConn sends from. A single address host sends everything from addr, whatever the
// destination. A dual stack host sends from the address matching the destination's family and has no route at all
// for a family it holds no address in.
type testerAddrs struct {
	addr   netip.AddrPort
	v4, v6 netip.AddrPort
	dual   bool
}

type TesterConn struct {
	// addrs is read by nebula's own goroutines on every send and by the router's flow renderer, and a test can
	// move it mid-run to simulate roaming, so it is atomic rather than a plain field.
	addrs atomic.Pointer[testerAddrs]

	RxPackets chan *Packet // Packets to receive into nebula
	TxPackets chan *Packet // Packets transmitted outside by nebula

	// done is closed exactly once by Close. Senders select on it so they
	// never race with a channel close; readers exit when it fires. The
	// packet channels are intentionally never closed - that was the source
	// of `send on closed channel` panics when a WriteTo/Send from another
	// goroutine passed the close check and reached the send just after
	// Close ran.
	done      chan struct{}
	closeOnce sync.Once

	l *slog.Logger
}

func NewListener(l *slog.Logger, s Settings) (Conn, error) {
	c := &TesterConn{
		RxPackets: make(chan *Packet, 10),
		TxPackets: make(chan *Packet, 10),
		done:      make(chan struct{}),
		l:         l,
	}
	c.SetAddr(s.Listen)
	return c, nil
}

// GetAddr returns the underlay address this conn currently sends from. For a dual stack conn that is the IPv4
// address, or the IPv6 one when there is no IPv4 service.
func (u *TesterConn) GetAddr() netip.AddrPort {
	return u.addrs.Load().addr
}

// SetAddr moves this conn to a new underlay address, standing in for a host waking up on a different network.
func (u *TesterConn) SetAddr(addr netip.AddrPort) {
	u.addrs.Store(&testerAddrs{addr: addr})
}

// SetDualStackAddrs gives this conn one underlay address per family, standing in for a host whose routing table picks
// the source by destination. Pass an invalid AddrPort for a family the network does not offer, sends to that family
// then fail with ENETUNREACH like they would on a host with no route for it.
func (u *TesterConn) SetDualStackAddrs(v4, v6 netip.AddrPort) {
	a := &testerAddrs{v4: v4, v6: v6, dual: true}
	if v4.IsValid() {
		a.addr = v4
	} else {
		a.addr = v6
	}
	u.addrs.Store(a)
}

// sourceFor returns the address a send to dst leaves from
func (u *TesterConn) sourceFor(dst netip.AddrPort) (netip.AddrPort, error) {
	a := u.addrs.Load()
	if !a.dual {
		return a.addr, nil
	}

	src := a.v6
	if dst.Addr().Unmap().Is4() {
		src = a.v4
	}
	if !src.IsValid() {
		return netip.AddrPort{}, &net.OpError{Op: "write", Net: "udp", Addr: net.UDPAddrFromAddrPort(dst), Err: syscall.ENETUNREACH}
	}
	return src, nil
}

// Send will place a UdpPacket onto the receive queue for nebula to consume
// this is an encrypted packet or a handshake message in most cases
// packets were transmitted from another nebula node, you can send them with Tun.Send
func (u *TesterConn) Send(packet *Packet) {
	if u.l.Enabled(context.Background(), slog.LevelDebug) {
		// Parse the header only under debug logging, otherwise the
		// allocation would show up in every Send call.
		var h header.H
		if err := h.Parse(packet.Data); err != nil {
			panic(err)
		}
		u.l.Debug("UDP receiving injected packet",
			"header", &h,
			"udpAddr", packet.From,
			"dataLen", len(packet.Data),
		)
	}
	select {
	case <-u.done:
	case u.RxPackets <- packet:
	}
}

// Get will pull a UdpPacket from the transmit queue
// nebula meant to send this message on the network, it will be encrypted
// packets were ingested from the tun side (in most cases), you can send them with Tun.Send
func (u *TesterConn) Get(block bool) *Packet {
	if block {
		select {
		case <-u.done:
			return nil
		case p := <-u.TxPackets:
			return p
		}
	}

	select {
	case p := <-u.TxPackets:
		return p
	default:
		return nil
	}
}

//********************************************************************************************************************//
// Below this is boilerplate implementation to make nebula actually work
//********************************************************************************************************************//

func (u *TesterConn) WriteTo(b []byte, addr netip.AddrPort) error {
	from, err := u.sourceFor(addr)
	if err != nil {
		return err
	}

	p := acquirePacket()
	if cap(p.Data) < len(b) {
		p.Data = make([]byte, len(b))
	} else {
		p.Data = p.Data[:len(b)]
	}
	copy(p.Data, b)
	p.From = from
	p.To = addr
	select {
	case <-u.done:
		p.Release()
		return io.ErrClosedPipe
	case u.TxPackets <- p:
		return nil
	}
}
func (u *TesterConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
	written := 0
	for i, b := range bufs {
		if err := u.WriteTo(b, addrs[i]); err == nil {
			written++
		} else {
			u.l.Debug("failed to write packet in batch", "udpAddr", addrs[i], "error", err)
		}
	}
	return written, nil
}

func (u *TesterConn) ListenOut(r EncReader, flush func()) error {
	for {
		select {
		case <-u.done:
			return os.ErrClosed
		case p := <-u.RxPackets:
			r(p.From, p.Data[:len(p.Data):len(p.Data)])
			// The batcher borrows plaintext decrypted in place inside p.Data
			// until Flush, so the packet must stay alive across flush()
			flush()
			p.Release()
		}
	}
}

func (u *TesterConn) ReloadConfig(*config.C, *config.Reload) {}

func NewUDPStatsEmitter(_ []Conn) func() {
	// No UDP stats for non-linux
	return func() {}
}

func (u *TesterConn) LocalAddr() (netip.AddrPort, error) {
	return u.GetAddr(), nil
}

func (u *TesterConn) SupportsMultipleReaders() bool {
	return false
}

func (u *TesterConn) Rebind() error {
	return nil
}

func (u *TesterConn) Close() error {
	u.closeOnce.Do(func() {
		close(u.done)
	})
	return nil
}