// Teleport
// Copyright (C) 2024 Gravitational, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package vnet

import (
	"context"
	"errors"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/gravitational/trace"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/singleflight"
	"golang.zx2c4.com/wireguard/device"
	"gvisor.dev/gvisor/pkg/buffer"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
	"gvisor.dev/gvisor/pkg/tcpip/header"
	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
	ipv4network "gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
	ipv6network "gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
	"gvisor.dev/gvisor/pkg/tcpip/stack"
	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
	"gvisor.dev/gvisor/pkg/waiter"

	"github.com/gravitational/teleport"
	logutils "github.com/gravitational/teleport/lib/utils/log"
	"github.com/gravitational/teleport/lib/vnet/dns"
)

var log = logutils.NewPackageLogger(teleport.ComponentKey, logComponent)

const (
	logComponent    = "vnet"
	dnsLogComponent = "dns"
	nicID           = 1
	// vnetTUNMTU is the default MTU for the TUN device. It can be overridden
	// with the TELEPORT_UNSTABLE_VNET_TUN_MTU environment variable.
	vnetTUNMTU = 16 * 1024
	// vnetTUNMTUEnvVar overrides the default TUN device MTU. It must be set in
	// the environment of the VNet admin process.
	vnetTUNMTUEnvVar = "TELEPORT_UNSTABLE_VNET_TUN_MTU"
	// minTUNMTU is the lowest accepted MTU override, chosen because IPv6
	// requires a minimum link MTU of 1280 bytes (RFC 8200).
	minTUNMTU = 1280
	// maxTUNMTU is the largest possible IP packet size, the 16-bit total
	// length field limit shared by IPv4 (RFC 791) and IPv6 (RFC 8200).
	maxTUNMTU                        = 65535
	tcpReceiveBufferSize             = 0 // 0 means a default will be used.
	maxInFlightTCPConnectionAttempts = 1024
)

// tunMTU returns the MTU to use for the TUN device.
func tunMTU(ctx context.Context) int {
	env := os.Getenv(vnetTUNMTUEnvVar)
	if env == "" {
		return vnetTUNMTU
	}
	mtu, err := strconv.Atoi(env)
	if err != nil || mtu < minTUNMTU || mtu > maxTUNMTU {
		log.WarnContext(ctx, "Ignoring invalid TUN MTU override.",
			"env_var", vnetTUNMTUEnvVar, "value", env, "min", minTUNMTU, "max", maxTUNMTU)
		return vnetTUNMTU
	}
	log.InfoContext(ctx, "Using TUN MTU override from environment.",
		"env_var", vnetTUNMTUEnvVar, "mtu", mtu)
	return mtu
}

// networkStackConfig holds configuration parameters for the VNet network stack.
type networkStackConfig struct {
	// tunDevice is the OS TUN virtual network interface.
	tunDevice TUNDevice
	// ipv6Prefix is the IPv6 ULA prefix to use for all assigned VNet IP addresses.
	ipv6Prefix tcpip.Address
	// dnsIPv6 is the IPv6 address on which to host the DNS server. It must be under IPv6Prefix.
	dnsIPv6 tcpip.Address
	// ipv6Disabled is true when IPv6 is disabled on the host. The network
	// stack keeps its internal IPv6 prefix, but VNet doesn't answer AAAA
	// queries, has no IPv6 DNS address, and doesn't register the IPv6
	// prefix on its TUN device.
	ipv6Disabled bool
	// tcpHandlerResolver will be used to resolve all DNS queries that VNet may
	// need to handle.
	tcpHandlerResolver *tcpHandlerResolver

	// upstreamNameserverSource, if set, overrides the default OS UpstreamNameserverSource which provides the
	// IP addresses that unmatched DNS queries should be forwarded to. It is used in tests.
	upstreamNameserverSource dns.UpstreamNameserverSource
}

// checkAndSetDefaults checks the config and sets defaults.
func (c *networkStackConfig) checkAndSetDefaults() error {
	if c.tunDevice == nil {
		return trace.BadParameter("tunDevice is required")
	}
	if c.ipv6Prefix.Len() != 16 || c.ipv6Prefix.AsSlice()[0] != 0xfd {
		return trace.BadParameter("ipv6Prefix must be an IPv6 ULA address")
	}
	if !c.ipv6Disabled && c.dnsIPv6 == (tcpip.Address{}) {
		return trace.BadParameter("dnsIPv6 is required when IPv6 is enabled")
	}
	if c.tcpHandlerResolver == nil {
		return trace.BadParameter("tcpHandlerResolver is required")
	}
	return nil
}

// getIPv6Prefix returns the IPv6 prefix to expose outside the network stack:
// empty when IPv6 is disabled on the host, even though ipv6Prefix is always
// set internally. Use it whenever the prefix leaves the network stack.
func (c *networkStackConfig) getIPv6Prefix() string {
	if c.ipv6Disabled {
		return ""
	}
	return c.ipv6Prefix.String()
}

// errNoTCPHandler should be returned by tcpHandlerResolvers when no handler
// matches the FQDN.
//
// Avoid using [trace.Wrap] on errNoTCPHandler where possible, this isn't an
// unexpected error that should require the overhead of collecting a stack trace.
var errNoTCPHandler = &trace.NotFoundError{Message: "no handler for address"}

// tcpHandlerSpec specifies a VNet TCP handler.
type tcpHandlerSpec struct {
	// ipv4CIDRRange is the network that any V4 IP address should be assigned to this handler from.
	ipv4CIDRRange string
	// tcpHandler is the handler for TCP connections.
	tcpHandler tcpHandler
}

// tcpHandler defines the behavior for handling TCP connections from VNet.
//
// Implementations should attempt to dial the target and return any errors before calling
// [connector] to complete the TCP handshake and get the TCP conn. This is so that clients will see that the
// TCP connection was refused, instead of seeing a successful TCP dial that is immediately closed.
type tcpHandler interface {
	handleTCPConnector(ctx context.Context, localPort uint16, connector func() (net.Conn, error)) error
}

// udpHandler defines the behavior for handling UDP connections from VNet.
type udpHandler interface {
	HandleUDP(context.Context, net.Conn) error
}

// TUNDevice abstracts a virtual network TUN device.
type TUNDevice interface {
	// Name returns the current name of the Device.
	Name() (string, error)

	// MTU returns the MTU of the Device. The network stack sizes its link
	// endpoint and packet buffers from this value, so it must not change over
	// the lifetime of a Device.
	MTU() (int, error)

	// Write one or more packets to the device (without any additional headers).
	// On a successful write it returns the number of packets written. A nonzero
	// offset can be used to instruct the Device on where to begin writing from
	// each packet contained within the bufs slice.
	Write(bufs [][]byte, offset int) (int, error)

	// Read one or more packets from the Device (without any additional headers).
	// On a successful read it returns the number of packets read, and sets
	// packet lengths within the sizes slice. len(sizes) must be >= len(bufs).
	// A nonzero offset can be used to instruct the Device on where to begin
	// reading into each element of the bufs slice.
	Read(bufs [][]byte, sizes []int, offset int) (n int, err error)

	// BatchSize returns the preferred/max number of packets that can be read or
	// written in a single read/write call. BatchSize must not change over the
	// lifetime of a Device.
	BatchSize() int

	// Close stops the Device and closes the Event channel.
	Close() error
}

// networkStack implements the TCP and UDP networking stack for VNet.
type networkStack struct {
	// stack is the gVisor networking stack.
	stack *stack.Stack

	// tun is the OS TUN device. Incoming IP/L3 packets will be copied from here to [linkEndpoint], and
	// outgoing packets from [linkEndpoint] will be written here.
	tun TUNDevice

	// linkEndpoint is the gVisor-side endpoint that emulates the OS TUN device. All incoming IP/L3 packets
	// from the OS TUN device will be injected as inbound packets to this endpoint to be processed by the
	// gVisor netstack which ultimately calls the TCP or UDP protocol handler. When the protocol handler
	// writes packets to the gVisor stack to an address assigned to this endpoint, they will be written to
	// this endpoint, and then copied from this endpoint to the OS TUN device.
	linkEndpoint *channel.Endpoint

	// ipv6Prefix holds the 96-bit prefix that will be used for all IPv6 addresses assigned in the VNet.
	ipv6Prefix tcpip.Address

	// ipv6Disabled is true when IPv6 is disabled on the host. The network
	// stack keeps its internal IPv6 prefix, but VNet doesn't answer AAAA
	// queries, has no IPv6 DNS address, and doesn't register the IPv6
	// prefix on its TUN device.
	ipv6Disabled bool

	// diagProbeIPv6 is the IPv6 address (ipv6Prefix::2) returned to diagnostic probe queries.
	// Set once in newNetworkStack.
	diagProbeIPv6 [16]byte

	// diagProbeIPv4 is the IPv4 address returned to diagnostic probe queries. Set on the
	// first call of addDNSAddress, which happens after the first call of targetOSConfig
	// and before DNS addresses are registered with any OS resolver, so ResolveA should
	// not read nil in practice.
	diagProbeIPv4 atomic.Pointer[[4]byte]

	// dnsServer is the VNet's local DNS server that can handle UDP DNS
	// requests.
	dnsServer *dns.Server
	// upstreamFilter removes VNet DNS addresses from upstream nameserver lists.
	upstreamFilter *filteredUpstreamSource

	// tcpHandlerResolver resolves FQDNs to a TCP handler that will be used to handle all future TCP
	// connections to IP addresses that will be assigned to that FQDN.
	tcpHandlerResolver *tcpHandlerResolver
	// resolveHandlerGroup is a [singleflight.Group] that will be used to avoid resolving the same FQDN
	// multiple times concurrently. Every call to [tcpHandlerResolver.ResolveTCPHandler] will be wrapped by
	// this. The key will be the FQDN.
	resolveHandlerGroup singleflight.Group

	// destroyed is a channel that will be closed when the VNet is in the process of being destroyed.
	// All goroutines should terminate quickly after either this is closed or the context passed to
	// [networkStack.Run] is canceled.
	destroyed chan struct{}
	// wg is a [sync.WaitGroup] that keeps track of all running goroutines started by the [networkStack].
	wg sync.WaitGroup

	// state holds all mutable state for the networkStack.
	state state

	slog *slog.Logger
}

type filteredUpstreamSource struct {
	base    dns.UpstreamNameserverSource
	mu      sync.RWMutex
	exclude map[string]struct{}
	slog    *slog.Logger
}

// filteredUpstreamSource wraps an upstream source and excludes addresses added via AddExclude.
// It is mainly used to filter VNet's own DNS addresses.
func newFilteredUpstreamSource(base dns.UpstreamNameserverSource, slog *slog.Logger) *filteredUpstreamSource {
	return &filteredUpstreamSource{
		base:    base,
		exclude: make(map[string]struct{}),
		slog:    slog,
	}
}

func (f *filteredUpstreamSource) AddExclude(addr string) {
	if addr == "" {
		return
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.exclude[addr] = struct{}{}
}

func (f *filteredUpstreamSource) UpstreamNameservers(ctx context.Context) ([]string, error) {
	nameservers, err := f.base.UpstreamNameservers(ctx)
	if err != nil {
		return nil, err
	}
	f.slog.DebugContext(ctx, "Loaded upstream nameservers (pre-filter)", "nameservers", nameservers)
	f.mu.RLock()
	defer f.mu.RUnlock()
	if len(f.exclude) == 0 {
		f.slog.DebugContext(ctx, "Loaded upstream nameservers (post-filter)", "nameservers", nameservers)
		return nameservers, nil
	}
	filtered := make([]string, 0, len(nameservers))
	for _, nameserver := range nameservers {
		if _, ok := f.exclude[nameserver]; ok {
			continue
		}
		filtered = append(filtered, nameserver)
	}
	f.slog.DebugContext(ctx, "Loaded upstream nameservers (post-filter)", "nameservers", filtered)
	return filtered, nil
}

type state struct {
	// mu is a single mutex that protects the whole state struct. This could be optimized as necessary.
	mu sync.RWMutex

	// Each FQDN gets assigned both an IPv4 address and an IPv6 address, where the 4-bit suffix of the IPv6
	// matches the IPv4 address exactly. All per-FQDN state references the smaller IPv4 address only and
	// lookups based on an IPv6 address can use the 4-byte suffix.

	// tcpHandlers holds the map of IP addresses to assigned TCP handlers.
	tcpHandlers map[ipv4]tcpHandler
	// assignedIPs holds the map of app FQDNs to their assigned IP address, it's a reverse map of [tcpHandlers].
	assignedIPs map[string]ipv4

	// udpHandlers holds the map of IP addresses to assigned UDP handlers.
	udpHandlers map[ipv4]udpHandler
}

func newState() state {
	return state{
		tcpHandlers: make(map[ipv4]tcpHandler),
		udpHandlers: make(map[ipv4]udpHandler),
		assignedIPs: make(map[string]ipv4),
	}
}

// newNetworkStack creates a new VNet network stack with the given configuration and root context.
// It takes ownership of [cfg.TUNDevice] and will handle closing it before Run() returns. Call Run()
// on the returned network stack to start the VNet.
func newNetworkStack(cfg *networkStackConfig) (*networkStack, error) {
	if err := cfg.checkAndSetDefaults(); err != nil {
		return nil, trace.Wrap(err)
	}
	logger := slog.With(teleport.ComponentKey, logComponent)
	dnsLogger := slog.With(teleport.ComponentKey, teleport.Component(logComponent, dnsLogComponent))

	mtu, err := cfg.tunDevice.MTU()
	if err != nil {
		return nil, trace.Wrap(err, "getting TUN device MTU")
	}
	stack, linkEndpoint, err := createStack(uint32(mtu))
	if err != nil {
		return nil, trace.Wrap(err)
	}

	if err := installVnetRoutes(stack); err != nil {
		return nil, trace.Wrap(err)
	}

	ns := &networkStack{
		tun:                cfg.tunDevice,
		stack:              stack,
		linkEndpoint:       linkEndpoint,
		ipv6Prefix:         cfg.ipv6Prefix,
		ipv6Disabled:       cfg.ipv6Disabled,
		tcpHandlerResolver: cfg.tcpHandlerResolver,
		destroyed:          make(chan struct{}),
		state:              newState(),
		slog:               logger,
	}

	upstreamNameserverSource := cfg.upstreamNameserverSource
	if upstreamNameserverSource == nil {
		upstreamNameserverSource, err = dns.NewOSUpstreamNameserverSource(dnsLogger)
		if err != nil {
			return nil, trace.Wrap(err)
		}
	}
	upstreamFilter := newFilteredUpstreamSource(upstreamNameserverSource, dnsLogger)
	ns.upstreamFilter = upstreamFilter
	if cfg.dnsIPv6 != (tcpip.Address{}) {
		addr, ok := netip.AddrFromSlice(cfg.dnsIPv6.AsSlice())
		if !ok {
			return nil, trace.Errorf("error parsing IPv6 DNS address %v", cfg.dnsIPv6)
		}
		upstreamFilter.AddExclude(dns.AddrWithDNSPort(addr))
	}
	dnsServer, err := dns.NewServer(ns, upstreamFilter, dnsLogger)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	ns.dnsServer = dnsServer

	tcpForwarder := tcp.NewForwarder(ns.stack, tcpReceiveBufferSize, maxInFlightTCPConnectionAttempts, ns.handleTCP)
	ns.stack.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)

	udpForwarder := udp.NewForwarder(ns.stack, ns.handleUDP)
	ns.stack.SetTransportProtocolHandler(udp.ProtocolNumber, udpForwarder.HandlePacket)

	if cfg.dnsIPv6 != (tcpip.Address{}) {
		ns.diagProbeIPv6 = cfg.dnsIPv6.As16()
		if err := ns.assignUDPHandler(cfg.dnsIPv6, dnsServer); err != nil {
			return nil, trace.Wrap(err)
		}
		logger.DebugContext(context.Background(), "Serving DNS on IPv6.", "dns_addr", cfg.dnsIPv6)
	} else {
		// Added it so the probe handler can always return a stable value
		ns.diagProbeIPv6 = ipv6WithSuffix(cfg.ipv6Prefix, dns.DNSServerSuffix).As16()
	}

	return ns, nil
}

func createStack(mtu uint32) (*stack.Stack, *channel.Endpoint, error) {
	netStack := stack.New(stack.Options{
		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4network.NewProtocol, ipv6network.NewProtocol},
		TransportProtocols: []stack.TransportProtocolFactory{tcp.NewProtocol, udp.NewProtocol},
	})

	const (
		size     = 512
		linkAddr = ""
	)
	linkEndpoint := channel.New(size, mtu, linkAddr)
	if err := netStack.CreateNIC(nicID, linkEndpoint); err != nil {
		return nil, nil, trace.Errorf("creating VNet NIC: %s", err)
	}

	return netStack, linkEndpoint, nil
}

func installVnetRoutes(stack *stack.Stack) error {
	// Make the network stack pass all outbound IP packets to the NIC, regardless of destination IP address.
	ipv4Subnet, err := tcpip.NewSubnet(tcpip.AddrFrom4([4]byte{}), tcpip.MaskFromBytes(make([]byte, 4)))
	if err != nil {
		return trace.Wrap(err, "creating VNet IPv4 subnet")
	}
	ipv6Subnet, err := tcpip.NewSubnet(tcpip.AddrFrom16([16]byte{}), tcpip.MaskFromBytes(make([]byte, 16)))
	if err != nil {
		return trace.Wrap(err, "creating VNet IPv6 subnet")
	}
	stack.SetRouteTable([]tcpip.Route{
		{
			Destination: ipv4Subnet,
			NIC:         nicID,
		},
		{
			Destination: ipv6Subnet,
			NIC:         nicID,
		},
	})
	return nil
}

// run starts the VNet networking stack. It blocks until [ctx] is canceled, at
// which point it closes the link endpoint, waits for all goroutines to
// terminate, and destroys the networking stack.
func (ns *networkStack) run(ctx context.Context) error {
	ns.slog.InfoContext(ctx, "Running Teleport VNet.", "ipv6_prefix", ns.ipv6Prefix)

	ctx, cancel := context.WithCancel(ctx)

	allErrors := make(chan error, 2)
	g, ctx := errgroup.WithContext(ctx)
	g.Go(func() error {
		// Make sure to cancel the context in case this exits prematurely with a nil error.
		defer cancel()
		err := forwardBetweenTunAndNetstack(ctx, ns.tun, ns.linkEndpoint)
		allErrors <- err
		return err
	})
	g.Go(func() error {
		// When the context is canceled for any reason (the caller or one of the other concurrent tasks may
		// have canceled it) destroy everything and quit.
		<-ctx.Done()
		ns.slog.InfoContext(ctx, "Context canceled, beginning network stack shutdown.")

		// In-flight connections should start terminating after closing [ns.destroyed].
		close(ns.destroyed)

		// Close the link endpoint and the TUN, this should cause [forwardBetweenTunAndNetstack] to terminate
		// if it hasn't already.
		ns.slog.InfoContext(ctx, "Closing link endpoint.")
		ns.linkEndpoint.Close()
		ns.slog.InfoContext(ctx, "Closing TUN device.")
		err := trace.Wrap(ns.tun.Close(), "closing TUN device")

		allErrors <- err
		return err
	})

	// Deliberately ignoring the error from g.Wait() to return an aggregate of all errors.
	_ = g.Wait()

	// Wait for all connections and goroutines to clean themselves up.
	ns.slog.InfoContext(ctx, "Waiting for all connections and goroutines to clean up.")
	ns.wg.Wait()

	// Now we can destroy the gVisor networking stack and wait for all its goroutines to terminate.
	ns.slog.InfoContext(ctx, "Destroying networking stack.")
	ns.stack.Destroy()

	close(allErrors)
	return trace.NewAggregateFromChannel(allErrors, context.Background())
}

func (ns *networkStack) handleTCP(req *tcp.ForwarderRequest) {
	// Add 1 to the waitgroup because the networking stack runs this in its own goroutine.
	ns.wg.Add(1)
	defer ns.wg.Done()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Clients of *tcp.ForwarderRequest must eventually call Complete on it exactly once.
	// [req] consumes 1 of [maxInFlightTCPConnectionAttempts] until [req.Complete] is called.
	var completed bool
	defer func() {
		if !completed {
			req.Complete(true /*send TCP reset*/)
		}
	}()

	id := req.ID()
	slog := ns.slog.With("request", id)
	slog.DebugContext(ctx, "Handling TCP connection.")
	defer slog.DebugContext(ctx, "Finished handling TCP connection.")

	handler, ok := ns.getTCPHandler(id.LocalAddress)
	if !ok {
		slog.DebugContext(ctx, "No handler for address.", "addr", id.LocalAddress)
		return
	}

	connector := func() (net.Conn, error) {
		var wq waiter.Queue
		waitEntry, notifyCh := waiter.NewChannelEntry(waiter.EventErr | waiter.EventHUp)
		wq.EventRegister(&waitEntry)

		endpoint, err := req.CreateEndpoint(&wq)
		if err != nil {
			// This err doesn't actually implement [error]
			return nil, trace.Errorf("creating TCP endpoint: %s", err)
		}

		completed = true
		req.Complete(false /*don't send TCP reset*/)

		endpoint.SocketOptions().SetKeepAlive(true)

		conn := gonet.NewTCPConn(&wq, endpoint)

		ns.wg.Add(1)
		go func() {
			defer func() {
				cancel()
				conn.Close()
				ns.wg.Done()
			}()
			select {
			case <-notifyCh:
				slog.DebugContext(ctx, "Got HUP or ERR, canceling request context and closing TCP conn.")
			case <-ns.destroyed:
				slog.DebugContext(ctx, "VNet is being destroyed, canceling request context and closing TCP conn.")
			case <-ctx.Done():
				slog.DebugContext(ctx, "Request context canceled, closing TCP conn.")
			}
		}()

		return conn, nil
	}

	if err := handler.handleTCPConnector(ctx, id.LocalPort, connector); err != nil {
		if errors.Is(err, context.Canceled) {
			slog.DebugContext(ctx, "TCP connection handler returned early due to canceled context.")
		} else {
			slog.DebugContext(ctx, "Error handling TCP connection.", "err", err)
		}
	}
}

func (ns *networkStack) getTCPHandler(addr tcpip.Address) (tcpHandler, bool) {
	ns.state.mu.RLock()
	defer ns.state.mu.RUnlock()
	handler, ok := ns.state.tcpHandlers[ipv4Suffix(addr)]
	return handler, ok
}

// assignTCPHandler assigns an IPv4 address to [handlerSpec] from its preferred CIDR range, and returns that
// new assigned address.
func (ns *networkStack) assignTCPHandler(handlerSpec *tcpHandlerSpec, fqdn string) (ipv4, error) {
	_, ipNet, err := net.ParseCIDR(handlerSpec.ipv4CIDRRange)
	if err != nil {
		return 0, trace.Wrap(err, "parsing CIDR %q", handlerSpec.ipv4CIDRRange)
	}

	ns.state.mu.Lock()
	defer ns.state.mu.Unlock()

	ip, err := randomFreeIPv4InNet(ipNet, func(ip ipv4) bool {
		_, taken := ns.state.tcpHandlers[ip]
		return !taken
	})
	if err != nil {
		return 0, trace.Wrap(err, "assigning IP address")
	}

	ns.state.tcpHandlers[ip] = handlerSpec.tcpHandler
	ns.state.assignedIPs[fqdn] = ip

	if err := ns.addProtocolAddress(tcpip.AddrFrom4(ip.asArray())); err != nil {
		return 0, trace.Wrap(err)
	}
	if err := ns.addProtocolAddress(ipv6WithSuffix(ns.ipv6Prefix, ip.asSlice())); err != nil {
		return 0, trace.Wrap(err)
	}

	return ip, nil
}

func (ns *networkStack) handleUDP(req *udp.ForwarderRequest) bool {
	ns.wg.Add(1)
	go func() {
		defer ns.wg.Done()
		ns.handleUDPConcurrent(req)
	}()
	return true
}

func (ns *networkStack) handleUDPConcurrent(req *udp.ForwarderRequest) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	id := req.ID()
	slog := ns.slog.With("request", id)
	slog.DebugContext(ctx, "Handling UDP request.")
	defer slog.DebugContext(ctx, "Finished handling UDP request.")

	handler, ok := ns.getUDPHandler(id.LocalAddress)
	if !ok {
		slog.DebugContext(ctx, "No handler for address.")
		return
	}

	var wq waiter.Queue
	waitEntry, notifyCh := waiter.NewChannelEntry(waiter.EventErr | waiter.EventHUp)
	wq.EventRegister(&waitEntry)

	endpoint, err := req.CreateEndpoint(&wq)
	if err != nil {
		slog.ErrorContext(ctx, "Failed to create UDP endpoint.", "error", err)
		return
	}

	conn := gonet.NewUDPConn(&wq, endpoint)
	defer conn.Close()

	ns.wg.Add(1)
	go func() {
		defer func() {
			cancel()
			conn.Close()
			ns.wg.Done()
		}()
		select {
		case <-notifyCh:
			slog.DebugContext(ctx, "Got HUP or ERR, canceling request context and closing UDP conn.")
		case <-ns.destroyed:
			slog.DebugContext(ctx, "VNet is being destroyed, canceling request context and closing UDP conn.")
		case <-ctx.Done():
			slog.DebugContext(ctx, "Request context canceled, closing UDP conn.")
		}
	}()

	if err := handler.HandleUDP(ctx, conn); err != nil {
		slog.DebugContext(ctx, "Error handling UDP conn.", "error", err)
	}
}

func (ns *networkStack) getUDPHandler(addr tcpip.Address) (udpHandler, bool) {
	ipv4 := ipv4Suffix(addr)
	ns.state.mu.RLock()
	defer ns.state.mu.RUnlock()
	handler, ok := ns.state.udpHandlers[ipv4]
	return handler, ok
}

func (ns *networkStack) assignUDPHandler(addr tcpip.Address, handler udpHandler) error {
	ipv4 := ipv4Suffix(addr)
	ns.state.mu.Lock()
	defer ns.state.mu.Unlock()
	if _, ok := ns.state.udpHandlers[ipv4]; ok {
		return trace.AlreadyExists("Handler for %s is already set", addr)
	}
	if err := ns.addProtocolAddress(addr); err != nil {
		return trace.Wrap(err)
	}
	ns.state.udpHandlers[ipv4] = handler
	return nil
}

// addDNSAddress adds a DNS handler at the given IP.
func (ns *networkStack) addDNSAddress(ip net.IP) error {
	slog.DebugContext(context.Background(), "Serving DNS on IPv4.", "dns_addr", ip.String())
	addr, ok := netip.AddrFromSlice(ip)
	if !ok {
		return trace.Errorf("error parsing IPv4 DNS address %s", ip.String())
	}
	if v4 := ip.To4(); v4 != nil && ns.diagProbeIPv4.Load() == nil {
		var b [4]byte
		copy(b[:], v4)
		ns.diagProbeIPv4.CompareAndSwap(nil, &b)
	}
	if ns.upstreamFilter != nil {
		ns.upstreamFilter.AddExclude(dns.AddrWithDNSPort(addr))
	}
	return trace.Wrap(ns.assignUDPHandler(tcpip.AddrFromSlice(ip), ns.dnsServer),
		"adding UDP handler at %s", ip.String())
}

// ResolveA implements [dns.Resolver.ResolveA].
func (ns *networkStack) ResolveA(ctx context.Context, fqdn string) (dns.Result, error) {
	// Diagnostic probes short-circuit here. Without this, each probe's unique random label
	// would allocate a fresh IPv4 from the CIDR pool, leaking the entry.
	if dns.HasDiagProbePrefix(fqdn) {
		if v4 := ns.diagProbeIPv4.Load(); v4 != nil {
			return dns.Result{A: *v4}, nil
		}
		return dns.Result{NoRecord: true}, nil
	}

	// Do the actual resolution within a [singleflight.Group] keyed by [fqdn] to avoid concurrent requests to
	// resolve an FQDN and then assign an address to it.
	resultAny, err, _ := ns.resolveHandlerGroup.Do(fqdn, func() (any, error) {
		// If we've already assigned an IP address to this fqdn, resolve to it.
		if ip, ok := ns.assignedIPv4(fqdn); ok {
			return dns.Result{
				A: ip.asArray(),
			}, nil
		}

		// If fqdn matches an address that should be handled by VNet, create a
		// new handler for it.
		handlerSpec, err := ns.tcpHandlerResolver.resolveTCPHandler(ctx, fqdn)
		if err != nil {
			if errors.Is(err, errNoTCPHandler) {
				// Did not find any match, forward the DNS request upstream.
				return dns.Result{}, nil
			}
			return dns.Result{}, trace.Wrap(err, "resolving TCP handler for fqdn %q", fqdn)
		}

		// Assign an unused IP address to this handler.
		ip, err := ns.assignTCPHandler(handlerSpec, fqdn)
		if err != nil {
			return dns.Result{}, trace.Wrap(err, "assigning address to handler for %q", fqdn)
		}

		// And resolve to the assigned address.
		return dns.Result{
			A: ip.asArray(),
		}, nil
	})
	if err != nil {
		return dns.Result{}, trace.Wrap(err)
	}
	return resultAny.(dns.Result), nil
}

// ResolveAAAA implements [dns.Resolver.ResolveAAAA].
func (ns *networkStack) ResolveAAAA(ctx context.Context, fqdn string) (dns.Result, error) {
	if ns.ipv6Disabled {
		// IPv6 is disabled on this host, return no AAAA records.
		return dns.Result{NoRecord: true}, nil
	}
	// Diagnostic probes return the stable IPv6 probe address — the value the diagnostic
	// check compares against. No handler is allocated.
	if dns.HasDiagProbePrefix(fqdn) {
		return dns.Result{AAAA: ns.diagProbeIPv6}, nil
	}

	result, err := ns.ResolveA(ctx, fqdn)
	if err != nil {
		return dns.Result{}, trace.Wrap(err)
	}
	if result.A != ([4]byte{}) {
		result.AAAA = ipv6WithSuffix(ns.ipv6Prefix, result.A[:]).As16()
		result.A = [4]byte{}
	}
	return result, nil
}

func (ns *networkStack) assignedIPv4(fqdn string) (ipv4, bool) {
	ns.state.mu.RLock()
	defer ns.state.mu.RUnlock()
	ipv4, ok := ns.state.assignedIPs[fqdn]
	return ipv4, ok
}

func forwardBetweenTunAndNetstack(ctx context.Context, tun TUNDevice, linkEndpoint *channel.Endpoint) error {
	slog.DebugContext(ctx, "Forwarding IP packets between OS and VNet.")
	g, ctx := errgroup.WithContext(ctx)
	g.Go(func() error { return forwardNetstackToTUN(ctx, linkEndpoint, tun) })
	g.Go(func() error { return forwardTUNtoNetstack(ctx, tun, linkEndpoint) })
	err := g.Wait()
	slog.DebugContext(ctx, "Finished forwarding IP packets between OS and VNet.")
	return trace.Wrap(err)
}

func forwardNetstackToTUN(ctx context.Context, linkEndpoint *channel.Endpoint, tun TUNDevice) error {
	mtu, err := tun.MTU()
	if err != nil {
		return trace.Wrap(err, "getting TUN device MTU")
	}
	bufs := [][]byte{make([]byte, device.MessageTransportHeaderSize+mtu)}
	for {
		packet := linkEndpoint.ReadContext(ctx)
		if packet == nil {
			// Nil packet is returned when context is canceled.
			return trace.Wrap(ctx.Err())
		}
		offset := device.MessageTransportHeaderSize
		views, pktOffset := packet.AsViewList()
		view := views.Front()
		for view != nil && pktOffset >= view.Size() {
			pktOffset -= view.Size()
			view = view.Next()
		}
		for ; view != nil; view = view.Next() {
			offset += copy(bufs[0][offset:], view.AsSlice()[pktOffset:])
			pktOffset = 0
		}
		packet.DecRef()
		bufs[0] = bufs[0][:offset]
		if _, err := tun.Write(bufs, device.MessageTransportHeaderSize); err != nil {
			return trace.Wrap(err, "writing packets to TUN")
		}
		bufs[0] = bufs[0][:cap(bufs[0])]
	}
}

// forwardTUNtoNetstack does not abort on ctx being canceled, but it does check the ctx error before
// returning os.ErrClosed from tun.Read.
func forwardTUNtoNetstack(ctx context.Context, tun TUNDevice, linkEndpoint *channel.Endpoint) error {
	const readOffset = device.MessageTransportHeaderSize
	mtu, err := tun.MTU()
	if err != nil {
		return trace.Wrap(err, "getting TUN device MTU")
	}
	bufs := make([][]byte, tun.BatchSize())
	for i := range bufs {
		bufs[i] = make([]byte, device.MessageTransportHeaderSize+mtu)
	}
	sizes := make([]int, len(bufs))
	for {
		n, err := tun.Read(bufs, sizes, readOffset)
		if err != nil {
			// tun.Read might get interrupted due to the TUN device getting closed after ctx cancellation.
			if errors.Is(err, os.ErrClosed) && ctx.Err() != nil {
				return ctx.Err()
			}
			return trace.Wrap(err, "reading packets from TUN")
		}
		for i := range sizes[:n] {
			protocol, ok := protocolVersion(bufs[i][readOffset])
			if !ok {
				continue
			}
			packet := stack.NewPacketBuffer(stack.PacketBufferOptions{
				Payload: buffer.MakeWithData(bufs[i][readOffset : readOffset+sizes[i]]),
			})
			linkEndpoint.InjectInbound(protocol, packet)
			packet.DecRef()
		}
	}
}

func (ns *networkStack) addProtocolAddress(localAddress tcpip.Address) error {
	protocolAddress, err := protocolAddress(localAddress)
	if err != nil {
		return trace.Wrap(err)
	}
	if err := ns.stack.AddProtocolAddress(nicID, protocolAddress, stack.AddressProperties{}); err != nil {
		return trace.Errorf("%s", err)
	}
	return nil
}

func protocolAddress(addr tcpip.Address) (tcpip.ProtocolAddress, error) {
	addrWithPrefix := addr.WithPrefix()
	var protocol tcpip.NetworkProtocolNumber
	switch addrWithPrefix.PrefixLen {
	case 32:
		protocol = ipv4network.ProtocolNumber
	case 128:
		protocol = ipv6network.ProtocolNumber
	default:
		return tcpip.ProtocolAddress{}, trace.BadParameter("unhandled prefix len %d", addrWithPrefix.PrefixLen)
	}
	return tcpip.ProtocolAddress{
		AddressWithPrefix: addrWithPrefix,
		Protocol:          protocol,
	}, nil
}

func protocolVersion(b byte) (tcpip.NetworkProtocolNumber, bool) {
	switch b >> 4 {
	case 4:
		return header.IPv4ProtocolNumber, true
	case 6:
		return header.IPv6ProtocolNumber, true
	}
	return 0, false
}