a73x

internal/agent/exposeproxy/exposeproxy.go

Ref:   Size: 16.0 KiB   History

// Package exposeproxy publishes guest ports on their host. A Manager owns one
// socket and one goroutine per exposure id — the same manager shape as the
// serial pumps — and Converge drives that set toward the exposures the fleet
// says this host should be serving: a new id binds, a vanished id closes, a
// changed spec closes and rebinds.
//
// A TCP exposure is a listener and a spliced connection per caller; a UDP one
// is a single bound socket and a session per client address (udp.go). The
// protocol is part of an exposure's spec, so changing it is a rebind like any
// other change.
//
// The socket binds 0.0.0.0. All interfaces is the point: the LAN address is
// the feature, the tailnet address is how an operator reaches it from
// elsewhere, and the guest-side bridge seeing it is the same trust domain the
// LAN posture already accepts. There is no authentication in front of a
// published port — reaching the host is reaching the service.
//
// Each accepted connection asks for the VM's address AT THAT MOMENT and dials
// the guest fresh. That per-connection lookup is what makes an exposure created
// before its guest has booted simply work: the socket binds now, and
// connections start succeeding when the guest does.
//
// Nothing here is persisted. Sockets are rebuilt from the first snapshot
// after the agent starts, the same way the consoles are; connections in flight
// across an agent restart drop, and reconnecting works.
//
// Every published port serves under a cap, because the agent that carries this
// proxy is also the thing that runs the guests: traffic on a published port
// must never be able to starve fleet management of file descriptors.
package exposeproxy

import (
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/a73x/eitri/internal/pb"
)

// maxConnsPerExposure is how many connections one published port may hold open
// at once. Each costs two descriptors — the caller's and the dial into the
// guest — so an exposure sitting at its cap spends 512 of the 65536 the shipped
// unit grants the agent: a host could serve dozens of published ports, all of
// them saturated, and still leave the agent's own descriptors (hypervisor
// children, disks, consoles, the sync tunnel) untouched. 256 at once is far
// more than a service behind a published port meets in real use, and a caller
// past it is closed immediately — a refusal the caller can see beats a proxy
// wedged in a way nobody can diagnose.
const maxConnsPerExposure = 256

// acceptRetryFloor and acceptRetryCeiling bound the pause an accept loop takes
// when the process is momentarily out of descriptors.
const (
	acceptRetryFloor   = 5 * time.Millisecond
	acceptRetryCeiling = time.Second
)

// Manager runs one socket per active exposure.
type Manager struct {
	// addr answers a VM's current guest address, or "" when this host does not
	// know one. Injected so the package stays a leaf: it never learns what a VM
	// is, only where one is right now. It is called concurrently from connection
	// goroutines; it must be safe for concurrent use.
	addr func(vmID string) string

	// maxConns, maxSessions and the two idle windows are the package constants,
	// held per-Manager so a test can serve the same behaviour with a handful of
	// callers instead of hundreds, and with windows it can outlast. Written
	// once at construction and only read after, so they need no lock.
	maxConns                   int64
	maxSessions                int
	unrepliedIdle, repliedIdle time.Duration

	mu   sync.Mutex
	live map[string]*exposure
}

// NewManager returns a Manager that resolves each connection's destination
// through addr.
func NewManager(addr func(vmID string) string) *Manager {
	return &Manager{
		addr:          addr,
		maxConns:      maxConnsPerExposure,
		maxSessions:   maxSessionsPerExposure,
		unrepliedIdle: udpUnrepliedIdle,
		repliedIdle:   udpRepliedIdle,
		live:          map[string]*exposure{},
	}
}

// exposure is one published port's running state: the spec key it is bound for,
// its socket — a listener for TCP, a packet socket for UDP, both nil when the
// bind failed — what the OS said if it did, and what it is holding open right
// now: connections for TCP, sessions for UDP.
//
// One entry lives per exposure id for as long as the fleet wants that id, and a
// changed spec closes and rebinds INSIDE it rather than replacing it. That is
// what keeps the count honest: connections spliced under the old spec drain
// through this same struct, and the descriptors they hold are still spent.
type exposure struct {
	key    string
	ln     net.Listener
	pc     *net.UDPConn
	reason string
	conns  atomic.Int64

	// refused and dropped are running totals for the life of this entry, which
	// is the life of the exposure on this agent: an id the fleet still wants
	// keeps its entry across rebinds, and an agent restart rebuilds every socket
	// from the first snapshot and starts both at zero. Totals rather than
	// gauges because what they count is momentary — see ExposureSessions in the
	// proto.
	refused atomic.Int64
	dropped atomic.Int64

	// sessions is the UDP session table, guarded by its own mutex because the
	// packet loop touches it on every datagram while the manager lock is held
	// across whole converges. nil until the first UDP bind.
	smu      sync.Mutex
	sessions map[string]*udpSession
}

// counters is what this exposure has carried, for the report. active is read
// from whichever half of the pipe this exposure is: a TCP exposure holds
// connections, a UDP one holds sessions, and no exposure is both.
func (e *exposure) counters() *pb.ExposureSessions {
	active := e.conns.Load()
	if e.pc != nil {
		e.smu.Lock()
		active = int64(len(e.sessions))
		e.smu.Unlock()
	}
	return &pb.ExposureSessions{
		Active:  active,
		Refused: e.refused.Load(),
		Dropped: e.dropped.Load(),
	}
}

// bound reports whether this exposure currently holds a socket.
func (e *exposure) bound() bool { return e.ln != nil || e.pc != nil }

// close tears down the socket, which is what ends the goroutine serving it, and
// drops every UDP session with it: the sessions belong to the socket they were
// created on, and there is nothing left to answer them through.
func (e *exposure) close() {
	if e.ln != nil {
		e.ln.Close()
		e.ln = nil
	}
	if e.pc != nil {
		e.pc.Close()
		e.pc = nil
	}
	e.closeSessions()
}

// specKey renders everything about a desired exposure that a bound socket
// depends on. A change to any of it is a close and a rebind rather than an
// edit — the socket's own address, or the protocol it speaks, is part of what
// changed.
func specKey(d *pb.ExposureSpec) string {
	return fmt.Sprintf("%s/%d/%d/%s", d.GetVmId(), d.GetGuestPort(), d.GetHostPort(), d.GetProtocol())
}

// Converge drives the running sockets toward desired and reports what each
// exposure is doing, in desired order. It is level-triggered: every call
// re-examines the whole set, so an exposure whose socket is gone — a failed
// bind, or a serving loop that died — is retried here and heals the moment the
// port frees.
//
// "active" means the socket is bound. Eitri owns the host half of the pipe;
// whether anything answers inside the guest is the guest's half, and this does
// not pretend otherwise. A bound exposure still carries a reason when something
// about serving it is going wrong — a descriptor shortage the accept loop is
// riding out — because a port that is bound and struggling is not the same
// thing as a port that is bound and fine.
func (m *Manager) Converge(desired []*pb.ExposureSpec) []*pb.ExposureStatus {
	m.mu.Lock()
	defer m.mu.Unlock()

	want := make(map[string]*pb.ExposureSpec, len(desired))
	for _, d := range desired {
		want[d.GetId()] = d
	}
	for id, ex := range m.live {
		if _, ok := want[id]; ok {
			continue
		}
		ex.close()
		delete(m.live, id)
	}

	out := make([]*pb.ExposureStatus, 0, len(desired))
	for _, d := range desired {
		ex, ok := m.live[d.GetId()]
		switch {
		case !ok:
			ex = &exposure{key: specKey(d)}
			m.live[d.GetId()] = ex
		case ex.key != specKey(d):
			// A rebind in place, not a fresh entry: the id is what the fleet
			// named and what callers are still draining through.
			ex.close()
			ex.key, ex.reason = specKey(d), ""
		}
		if !ex.bound() {
			if err := m.bind(d, ex); err != nil {
				ex.reason = err.Error()
				out = append(out, &pb.ExposureStatus{Id: d.GetId(), State: "failed", Reason: ex.reason, Sessions: ex.counters()})
				continue
			}
		}
		// Where the guest is can change under a live session — a VM re-imaged,
		// re-addressed, or gone. A session pins the address it was created for,
		// so this is where a pin that no longer matches is dropped: the next
		// datagram from that client starts a session pointing at wherever the
		// guest is now. Only UDP exposures ask, because only they hold anything
		// pinned, and asking is a read of the VM's record on disk.
		if ex.pc != nil {
			ex.evictMovedSessions(m.addr(d.GetVmId()))
		}
		out = append(out, &pb.ExposureStatus{Id: d.GetId(), State: "active", Reason: ex.reason, Sessions: ex.counters()})
	}
	return out
}

// bind opens the socket one desired exposure calls for and starts the loop that
// serves it. UDP is a bound packet socket; everything else is a TCP listener —
// a server that names no protocol at all means the one this proxy started with.
func (m *Manager) bind(d *pb.ExposureSpec, ex *exposure) error {
	hostAddr := net.JoinHostPort("0.0.0.0", strconv.Itoa(int(d.GetHostPort())))
	if d.GetProtocol() == "udp" {
		ua, err := net.ResolveUDPAddr("udp4", hostAddr)
		if err != nil {
			return err
		}
		pc, err := net.ListenUDP("udp4", ua)
		if err != nil {
			return err
		}
		ex.pc, ex.reason = pc, ""
		go m.serve(d.GetId(), ex, pc, d.GetVmId(), d.GetGuestPort())
		return nil
	}
	ln, err := net.Listen("tcp", hostAddr)
	if err != nil {
		return err
	}
	ex.ln, ex.reason = ln, ""
	go m.accept(d.GetId(), ex, ln, d.GetVmId(), d.GetGuestPort())
	return nil
}

// StopAll closes every socket (agent shutdown, tests).
func (m *Manager) StopAll() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for id, ex := range m.live {
		ex.close()
		delete(m.live, id)
	}
}

// accept serves one exposure's listener until it ends: a converge dropped it,
// StopAll closed it, or the listener itself broke. That third exit is why this
// gives the exposure back its listener-less state instead of just returning —
// an accept loop that died leaves a port bound that nothing is serving, and a
// caller would hang in the backlog while Converge went on reporting "active".
// Releasing the listener and recording why makes the next converge rebind it.
//
// The pointer-identity guard is what keeps the normal exits quiet: on a
// close-driven exit the entry is already gone, already nil, or already holds a
// newer listener, so only the loop that owns the current listener writes here.
//
// A descriptor shortage is the one accept error this does not die on. Handing
// the listener back would unbind a working port and ask the next converge to
// find a descriptor for a fresh one — exactly what the process has none of —
// so the loop pauses and keeps the port, and the backlog keeps callers waiting
// rather than refusing them. It says so in the exposure's reason for as long as
// the streak lasts: the port is bound, which is what "active" means, and an
// operator looking at a published port nothing is getting through deserves the
// one sentence that explains it. Everything else is a broken listener and heals
// the only way a broken listener can, by being rebound.
func (m *Manager) accept(id string, ex *exposure, ln net.Listener, vmID string, guestPort uint32) {
	var pause time.Duration
	capped := false
	for {
		conn, err := ln.Accept()
		if err != nil {
			if outOfDescriptors(err) {
				if pause == 0 {
					pause = acceptRetryFloor
					slog.Warn("exposure accept out of descriptors, retrying", "exposure", id, "err", err)
					m.setReason(id, ln, "accept: out of descriptors, retrying: "+err.Error())
				} else {
					pause = min(pause*2, acceptRetryCeiling)
				}
				time.Sleep(pause)
				continue
			}
			m.mu.Lock()
			if cur, ok := m.live[id]; ok && cur.ln == ln {
				cur.close()
				cur.reason = "accept: " + err.Error()
				// The reason travels in the next report, which nobody watching
				// the host sees; a broken listener belongs in the agent's log
				// too.
				slog.Warn("exposure accept loop died", "exposure", id, "err", err)
			}
			m.mu.Unlock()
			return
		}
		if pause != 0 {
			pause = 0
			slog.Info("exposure accept has descriptors again", "exposure", id)
			m.setReason(id, ln, "")
		}

		// Increment first and give it back if the cap says no: counting only
		// after the decision would let a burst of simultaneous accepts each read
		// a stale count and all pass.
		if ex.conns.Add(1) > m.maxConns {
			ex.conns.Add(-1)
			ex.refused.Add(1)
			conn.Close()
			slog.Debug("exposure refused a connection at its cap", "exposure", id, "cap", m.maxConns, "client", conn.RemoteAddr())
			if !capped {
				capped = true
				slog.Warn("exposure at its connection cap, refusing", "exposure", id, "cap", m.maxConns)
			}
			continue
		}
		capped = false
		slog.Debug("exposure connection opened", "exposure", id, "client", conn.RemoteAddr())
		go func() {
			defer ex.conns.Add(-1)
			m.pipe(ex, conn, vmID, guestPort)
			slog.Debug("exposure connection closed", "exposure", id, "client", conn.RemoteAddr())
		}()
	}
}

// setReason records what an accept loop wants the next report to say about a
// still-bound exposure. It writes under the manager lock, which is what makes
// the reason a converge reads a whole one rather than a torn one, and only when
// the entry still holds THIS loop's listener — the same identity guard the
// loop's exit takes, because a loop whose exposure was rebound underneath it
// has nothing to say about the listener that replaced its own.
func (m *Manager) setReason(id string, ln net.Listener, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if cur, ok := m.live[id]; ok && cur.ln == ln {
		cur.reason = reason
	}
}

// outOfDescriptors reports whether an accept failed because the process (EMFILE)
// or the machine (ENFILE) has no descriptor to give it. The runtime already
// retries the other transient accept errors — an interrupted call, a caller that
// went away between SYN and accept — so those never reach here.
func outOfDescriptors(err error) bool {
	return errors.Is(err, syscall.EMFILE) || errors.Is(err, syscall.ENFILE)
}

// pipe connects one accepted connection to the guest. No address (a guest still
// leasing) or a guest that will not answer closes immediately: the host half of
// the pipe exists, the guest half does not, and waiting would only hold the
// caller open on a promise nothing is keeping. Both are counted as drops — from
// the caller's side the port accepted and then said nothing, and the exposure's
// own totals are the only place that difference is written down.
func (m *Manager) pipe(ex *exposure, client net.Conn, vmID string, guestPort uint32) {
	ip := m.addr(vmID)
	if ip == "" {
		ex.dropped.Add(1)
		client.Close()
		return
	}
	guest, err := net.Dial("tcp", net.JoinHostPort(ip, strconv.Itoa(int(guestPort))))
	if err != nil {
		ex.dropped.Add(1)
		client.Close()
		return
	}
	splice(client, guest)
}

// splice copies in both directions until each ends, closing both connections
// once they have.
func splice(a, b net.Conn) {
	defer a.Close()
	defer b.Close()
	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); copyThenCloseWrite(b, a) }()
	go func() { defer wg.Done(); copyThenCloseWrite(a, b) }()
	wg.Wait()
}

// copyThenCloseWrite copies src into dst and then shuts down dst's WRITE half
// only, so the peer sees a clean EOF while the opposite direction keeps
// flowing. Half-close is what makes a protocol that signals end-of-request by
// closing its write side work through the proxy instead of hanging. A
// connection with no half-close is closed outright — consumer-side interface,
// so nothing here depends on the concrete net type.
func copyThenCloseWrite(dst, src net.Conn) {
	_, _ = io.Copy(dst, src)
	if cw, ok := dst.(interface{ CloseWrite() error }); ok {
		_ = cw.CloseWrite()
		return
	}
	_ = dst.Close()
}