a73x

internal/agent/exposeproxy/udp.go

Ref:   Size: 10.6 KiB   History

package exposeproxy

import (
	"errors"
	"log/slog"
	"net"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

// A UDP exposure is one bound socket on the host and one session per client
// address behind it. The session is the whole design: UDP has no connection to
// follow, so the proxy invents the smallest thing that lets a reply find its
// way home — a connected socket toward the guest, whose local port is what the
// guest answers to, held for as long as the conversation looks live.
//
// The guest's address is resolved ONCE, when the session is created, and pinned
// there. A datagram that arrives before this host knows where the guest is is
// dropped and starts nothing: a session pinned to nowhere would have to
// re-resolve on some later packet, and a proxy that re-resolves mid-conversation
// is one that can quietly start sending a guest's traffic to whatever took its
// address over. Converge is where a stale pin is dropped instead.
const (
	// maxSessionsPerExposure is how many client conversations one published UDP
	// port may hold at once. Each is one descriptor — the socket toward the
	// guest; the socket callers reach is shared by all of them — so the sizing
	// rationale is the TCP cap's, one descriptor per session instead of two,
	// against the same 65536 the shipped unit grants the agent. A table at its
	// cap refuses new clients and keeps the ones it has: eviction would trade a
	// conversation that is working for one that might not be, and a UDP client
	// has no close to notice.
	maxSessionsPerExposure = 256

	// udpUnrepliedIdle and udpRepliedIdle are how long a session outlives its
	// last datagram. A session the guest has never answered is a guess — a
	// scanner's single packet, or a service that was not listening — and thirty
	// seconds is long enough for a slow first reply and short enough that a
	// sweep of the port space does not fill the table. Once traffic has flowed
	// back the session is a real conversation, and two minutes carries the
	// quiet stretches those have (a DNS client between queries, a game tick
	// paused, an NTP poll interval).
	udpUnrepliedIdle = 30 * time.Second
	udpRepliedIdle   = 120 * time.Second

	// udpDatagramMax is the read buffer on both halves — larger than the
	// largest datagram either side can send, so what goes in one end comes out
	// the other whole. One datagram in is one datagram out: nothing here
	// batches, coalesces, or reassembles.
	udpDatagramMax = 64 * 1024
)

// udpSession is one client address's conversation with one guest port: the
// connected socket the guest sees, the client to send its answers back to, and
// the two facts that decide when the session ends.
type udpSession struct {
	// key is the client address as the session table indexes it.
	key    string
	client *net.UDPAddr

	// guest is CONNECTED to the pinned address, which is what makes a reply
	// readable here rather than on the shared socket — and what surfaces an
	// ICMP port-unreachable as a read error instead of silence.
	guest *net.UDPConn

	// pinned is the guest address this session was created for. Converge
	// compares it against where the VM is now.
	pinned string

	// unrepliedIdle and repliedIdle are the manager's two windows, copied here
	// so a session decides its own expiry without reaching back for them.
	unrepliedIdle, repliedIdle time.Duration

	// last is the unix-nano time of the last datagram in either direction, and
	// replied records that the guest has answered at least once. Both are
	// touched by the packet loop and the relay goroutine, so both are atomic.
	last    atomic.Int64
	replied atomic.Bool

	closeOnce sync.Once
}

func newUDPSession(client *net.UDPAddr, guest *net.UDPConn, pinned string, unreplied, replied time.Duration) *udpSession {
	s := &udpSession{
		key: client.String(), client: client, guest: guest, pinned: pinned,
		unrepliedIdle: unreplied, repliedIdle: replied,
	}
	s.touch()
	return s
}

// touch records that a datagram just moved.
func (s *udpSession) touch() { s.last.Store(time.Now().UnixNano()) }

// expiry is when the session ends if nothing else moves: its last datagram plus
// the window its promotion earns it.
func (s *udpSession) expiry() time.Time {
	idle := s.unrepliedIdle
	if s.replied.Load() {
		idle = s.repliedIdle
	}
	return time.Unix(0, s.last.Load()).Add(idle)
}

// close releases the session's descriptor, which is also what ends its relay.
func (s *udpSession) close() {
	s.closeOnce.Do(func() { s.guest.Close() })
}

// session returns the session serving client, or nil when there is none.
func (e *exposure) session(key string) *udpSession {
	e.smu.Lock()
	defer e.smu.Unlock()
	return e.sessions[key]
}

// addSession puts s in the table unless it is full, in which case the caller
// closes s and the datagram that would have started it is dropped.
func (e *exposure) addSession(s *udpSession, max int) bool {
	e.smu.Lock()
	defer e.smu.Unlock()
	if len(e.sessions) >= max {
		return false
	}
	if e.sessions == nil {
		e.sessions = map[string]*udpSession{}
	}
	e.sessions[s.key] = s
	return true
}

// dropSession removes s and closes it. Idempotent, and identity-checked: a
// session that has already been replaced under its key is not the one to
// remove.
func (e *exposure) dropSession(s *udpSession) {
	e.smu.Lock()
	if cur, ok := e.sessions[s.key]; ok && cur == s {
		delete(e.sessions, s.key)
	}
	e.smu.Unlock()
	s.close()
}

// closeSessions empties the table, ending every relay with it.
func (e *exposure) closeSessions() {
	e.smu.Lock()
	held := e.sessions
	e.sessions = nil
	e.smu.Unlock()
	for _, s := range held {
		s.close()
	}
}

// evictMovedSessions drops every session whose pinned guest address is not
// where the VM is now — including when the host no longer knows where that is,
// which is not an address to keep sending to either.
func (e *exposure) evictMovedSessions(guestIP string) {
	e.smu.Lock()
	var moved []*udpSession
	for key, s := range e.sessions {
		if s.pinned != guestIP {
			delete(e.sessions, key)
			moved = append(moved, s)
		}
	}
	e.smu.Unlock()
	for _, s := range moved {
		s.close()
	}
}

// serve reads one exposure's published UDP socket until it ends: a converge
// dropped it, StopAll closed it, or the socket itself broke. Like the accept
// loop, a broken socket is handed back so the next converge rebinds it — a
// bound port nothing is reading is worse than no port at all, because it
// swallows datagrams instead of refusing them.
//
// Every datagram either continues a session or starts one. Nothing here waits:
// a client with nowhere to send, a guest that will not take a connected socket,
// and a table at its cap all drop the datagram, because a datagram held is a
// datagram late, and UDP callers retry.
func (m *Manager) serve(id string, ex *exposure, pc *net.UDPConn, vmID string, guestPort uint32) {
	buf := make([]byte, udpDatagramMax)
	full := false
	for {
		n, client, err := pc.ReadFromUDP(buf)
		if err != nil {
			m.mu.Lock()
			if cur, ok := m.live[id]; ok && cur.pc == pc {
				cur.close()
				cur.reason = "read: " + err.Error()
				// The reason travels in the next report, which nobody watching
				// the host sees; a broken socket belongs in the agent's log too.
				slog.Warn("exposure packet loop died", "exposure", id, "err", err)
			}
			m.mu.Unlock()
			return
		}

		s := ex.session(client.String())
		if s == nil {
			s = m.openSession(id, ex, pc, client, vmID, guestPort, &full)
			if s == nil {
				continue
			}
		}
		if _, err := s.guest.Write(buf[:n]); err != nil {
			ex.dropped.Add(1)
			slog.Debug("exposure session torn down, send to guest failed", "exposure", id, "client", s.key, "err", err)
			ex.dropSession(s)
			continue
		}
		s.touch()
	}
}

// openSession resolves where the guest is right now, pins a connected socket
// there, and files it under the client's address. It returns nil when the
// datagram that asked for it is to be dropped — no guest address, a socket the
// OS would not give, or a table with no room — with full carrying the
// at-capacity streak so a saturated port logs once rather than per packet.
//
// The two ways of returning nil are counted apart, because they are different
// answers to an operator's question. A table at its cap is this proxy's own
// limit and says raise it; everything else here is the guest not being
// reachable yet, and says look at the guest.
func (m *Manager) openSession(id string, ex *exposure, pc *net.UDPConn, client *net.UDPAddr, vmID string, guestPort uint32, full *bool) *udpSession {
	ip := m.addr(vmID)
	if ip == "" {
		ex.dropped.Add(1)
		return nil
	}
	guestAddr, err := net.ResolveUDPAddr("udp4", net.JoinHostPort(ip, strconv.Itoa(int(guestPort))))
	if err != nil {
		ex.dropped.Add(1)
		return nil
	}
	guest, err := net.DialUDP("udp4", nil, guestAddr)
	if err != nil {
		ex.dropped.Add(1)
		return nil
	}
	s := newUDPSession(client, guest, ip, m.unrepliedIdle, m.repliedIdle)
	if !ex.addSession(s, m.maxSessions) {
		s.close()
		ex.refused.Add(1)
		slog.Debug("exposure refused a session at its cap", "exposure", id, "cap", m.maxSessions, "client", client.String())
		if !*full {
			*full = true
			slog.Warn("exposure at its UDP session cap, dropping", "cap", m.maxSessions, "client", client.String())
		}
		return nil
	}
	*full = false
	slog.Debug("exposure session opened", "exposure", id, "client", client.String(), "guest", guestAddr.String())
	go m.relay(id, ex, pc, s)
	return s
}

// relay carries one session's replies back to its client until the session
// ends. It ends on the guest socket going quiet for longer than the session's
// window, on any error reading it — which is where an ICMP port-unreachable
// from a guest with nothing listening arrives — and on a failure to answer the
// client.
//
// The read deadline IS the idle timer: there is no sweeper, because the only
// thing that needs to notice an expired session is the goroutine already
// waiting on it. A deadline that fires early because the forward direction
// moved the session on is simply re-armed against the newer expiry.
func (m *Manager) relay(id string, ex *exposure, pc *net.UDPConn, s *udpSession) {
	defer func() {
		ex.dropSession(s)
		slog.Debug("exposure session closed", "exposure", id, "client", s.key)
	}()
	buf := make([]byte, udpDatagramMax)
	for {
		if err := s.guest.SetReadDeadline(s.expiry()); err != nil {
			return
		}
		n, err := s.guest.Read(buf)
		if err != nil {
			if errors.Is(err, os.ErrDeadlineExceeded) && time.Now().Before(s.expiry()) {
				continue
			}
			return
		}
		if _, err := pc.WriteToUDP(buf[:n], s.client); err != nil {
			return
		}
		s.touch()
		// Promotion, and it only happens here: traffic that has come BACK from
		// the guest is what tells a real conversation apart from a stray packet
		// at a port.
		s.replied.Store(true)
	}
}