flow_dispatch.go

package tun

import (
	"bytes"
	"encoding/binary"
	"maps"
	"net/netip"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sagernet/sing-tun/gtcpip"
	"github.com/sagernet/sing-tun/gtcpip/checksum"
	"github.com/sagernet/sing-tun/gtcpip/header"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/logger"
)

const (
	tcpEstablishedTimeout = 2*time.Hour + 4*time.Minute
	tcpTransitoryTimeout  = 4 * time.Minute
	tcpClosingTimeout     = 10 * time.Second

	defaultUDPTimeout = 5 * time.Minute

	defaultICMPTimeout = time.Minute

	flowTombstoneTimeout = 4 * time.Minute

	flowTableCapacity = 16384
	flowStageCapacity = 1024

	flowSweepInterval = 30 * time.Second
	flowSweepLimit    = flowTableCapacity / int(flowTombstoneTimeout/flowSweepInterval)

	flowHoldCapacity = 512
	flowHoldPerFlow  = 16
)

type ForwardWriteback interface {
	ReturnHeadroom() int
	WriteReturnPackets(packets [][]byte) error
}

type ForwardFrameMeta struct {
	needsChecksum  bool
	checksumStart  uint16
	checksumOffset uint16
	gsoType        uint8
	gsoSize        uint16
}

func (m *ForwardFrameMeta) completeChecksum(raw []byte) {
	if !m.needsChecksum {
		return
	}
	m.needsChecksum = false
	checksumAt := int(m.checksumStart) + int(m.checksumOffset)
	if int(m.checksumStart) >= len(raw) || checksumAt+2 > len(raw) {
		return
	}
	initial := binary.BigEndian.Uint16(raw[checksumAt:])
	raw[checksumAt], raw[checksumAt+1] = 0, 0
	transportChecksum := ^checksum.Checksum(raw[m.checksumStart:], initial)
	if transportChecksum == 0 {
		transportChecksum = 0xffff
	}
	binary.BigEndian.PutUint16(raw[checksumAt:], transportChecksum)
}

type flowEntry struct {
	action   FlowAction
	deadline int64
	idle     time.Duration
	flow     *forwardFlow
	verdict  FlowVerdict
}

type forwardFlow struct {
	nat          *portNAT
	mapping      *natMapping
	selectorKey  natSelectorKey
	owner        *ForwardStage
	reverseKey   flowKey
	forwardRule  rewriteRule
	reverseRule  rewriteRule
	effectiveMTU uint32
	protocol     uint8
	udpTimeout   time.Duration
	tracker      FlowTracker

	clientAddress            netip.Addr
	clientSelector           uint16
	clientDestinationAddress netip.Addr
	clientDestinationPort    uint16
	serverAddress            netip.Addr
	dnatAddress              bool
	dnatPort                 bool

	finForward  atomic.Bool
	established atomic.Bool
	finReverse  atomic.Bool
	reported    atomic.Bool
	closed      atomic.Bool
	aborted     atomic.Bool
	lastReverse atomic.Int64
}

func (f *forwardFlow) report(reason FlowCloseReason) {
	if !f.reported.CompareAndSwap(false, true) {
		return
	}
	if f.tracker != nil {
		f.tracker.CloseFlow(reason)
	}
}

func (f *forwardFlow) close(reason FlowCloseReason) {
	if !f.closed.CompareAndSwap(false, true) {
		return
	}
	f.report(reason)
}

func (f *forwardFlow) CloseFlow() {
	f.aborted.Store(true)
	f.close(FlowCloseReset)
}

func (f *forwardFlow) observeReverse(packet *forwardPacket, now int64) {
	f.lastReverse.Store(now)
	if packet.protocol != uint8(header.TCPProtocolNumber) {
		return
	}
	if f.established.CompareAndSwap(false, true) && f.tracker != nil {
		f.tracker.FlowEstablished()
	}
	if packet.tcpFlags&header.TCPFlagRst != 0 {
		f.close(FlowCloseReset)
		return
	}
	if packet.tcpFlags&header.TCPFlagFin != 0 {
		f.finReverse.Store(true)
	} else if f.finReverse.Load() && f.finForward.Load() {
		f.report(FlowCloseFinished)
	}
}

type ForwardDispatcher struct {
	epoch        time.Time
	handler      Handler
	writeback    ForwardWriteback
	logger       logger.Logger
	mtu          uint32
	udpTimeout   time.Duration
	udpMapping   NATMapping
	udpFiltering NATFiltering
	icmpTimeout  time.Duration

	closed          atomic.Bool
	portsAccess     sync.Mutex
	ports           map[Port]*portReturn
	returnsAccess   sync.Mutex
	returns         atomic.Pointer[[]*portReturn]
	stagesAccess    sync.Mutex
	stages          []*ForwardStage
	stageCount      atomic.Int32
	resetGeneration atomic.Uint32
	identification  atomic.Uint32
}

type stagedPort struct {
	nat     *portNAT
	packets [][]byte
}

type heldPacket struct {
	packet          []byte
	transportLength int
	meta            ForwardFrameMeta
}

type flowJudgement struct {
	key     flowKey
	packets []heldPacket
	verdict FlowVerdict
}

type ForwardStage struct {
	dispatcher *ForwardDispatcher
	writeback  ForwardWriteback
	resume     func()
	access     sync.Mutex
	table      map[flowKey]*flowEntry
	judging    map[flowKey]*flowJudgement
	judged     []*flowJudgement
	held       int
	lastSweep  int64
	resetSeen  uint32

	exhaustedLogAt int64
	ports          []stagedPort
	returns        []returnBatch
	portOutput     stagePortOutput
	returnOutput   stageReturnOutput
	writebackBatch [][]byte
	segmentBuffers [][]byte
	segmentSizes   []int
	segmentUsed    int
}

func addrToTCPIP(addr netip.Addr) tcpip.Address {
	if addr.Is4() {
		return tcpip.AddrFrom4(addr.As4())
	}
	return tcpip.AddrFrom16(addr.As16())
}

func NewForwardDispatcher(handler Handler, writeback ForwardWriteback, logger logger.Logger, mtu uint32, udpOptions UDPNatOptions, icmpTimeout time.Duration) *ForwardDispatcher {
	dispatcher := &ForwardDispatcher{
		epoch:        time.Now(),
		handler:      handler,
		writeback:    writeback,
		logger:       logger,
		mtu:          mtu,
		udpTimeout:   udpOptions.Timeout,
		udpMapping:   udpOptions.Mapping,
		udpFiltering: udpOptions.Filtering,
		icmpTimeout:  icmpTimeout,
		ports:        make(map[Port]*portReturn),
	}
	if dispatcher.udpTimeout <= 0 {
		dispatcher.udpTimeout = defaultUDPTimeout
	}
	if dispatcher.icmpTimeout <= 0 {
		dispatcher.icmpTimeout = defaultICMPTimeout
	}
	dispatcher.returns.Store(new([]*portReturn))
	return dispatcher
}

func (d *ForwardDispatcher) NewStage(writeback ForwardWriteback, resume func()) *ForwardStage {
	if d == nil {
		return nil
	}
	if writeback == nil {
		writeback = d.writeback
	}
	stage := &ForwardStage{
		dispatcher: d,
		writeback:  writeback,
		resume:     resume,
		table:      make(map[flowKey]*flowEntry),
		judging:    make(map[flowKey]*flowJudgement),
		resetSeen:  d.resetGeneration.Load(),
	}
	stage.portOutput.stage = stage
	stage.returnOutput.stage = stage
	d.stagesAccess.Lock()
	d.stages = append(d.stages, stage)
	d.stageCount.Store(int32(len(d.stages)))
	d.stagesAccess.Unlock()
	return stage
}

func (d *ForwardDispatcher) now() int64 {
	return int64(time.Since(d.epoch))
}

func (d *ForwardDispatcher) Close() {
	if d == nil {
		return
	}
	d.closed.Store(true)
	d.stagesAccess.Lock()
	stages := d.stages
	d.stagesAccess.Unlock()
	for _, stage := range stages {
		stage.access.Lock()
		for key, entry := range stage.table {
			stage.removeEntry(key, entry, FlowCloseReset)
		}
		clear(stage.judging)
		stage.judged = nil
		stage.held = 0
		stage.access.Unlock()
	}
	d.portsAccess.Lock()
	returnPaths := slices.Collect(maps.Values(d.ports))
	d.portsAccess.Unlock()
	for _, returnPath := range returnPaths {
		returnPath.nat.port.DetachReturn(returnPath)
		releasePortNAT(returnPath.nat)
	}
}

func (s *ForwardStage) Dispatch(packet []byte) bool {
	if s == nil || s.dispatcher.closed.Load() {
		return false
	}
	parsed, ok := parseForwardPacket(packet)
	if !ok {
		return false
	}
	return s.dispatch(packet, &parsed, nil)
}

func (s *ForwardStage) DispatchParsed(packet []byte, meta ForwardFrameMeta, parsed *forwardPacket) bool {
	if s == nil || s.dispatcher.closed.Load() {
		return false
	}
	return s.dispatch(packet, parsed, &meta)
}

func (s *ForwardStage) dispatch(packet []byte, parsed *forwardPacket, meta *ForwardFrameMeta) bool {
	if parsed.fragment || !parsed.hasFlow {
		return false
	}
	key := parsed.flowKey()
	now := s.dispatcher.now()
	s.access.Lock()
	defer s.access.Unlock()
	entry, loaded := s.table[key]
	if loaded {
		if entryExpired(entry, now) {
			s.removeEntry(key, entry, FlowCloseTimeout)
		} else {
			handled, remove := s.handleHit(entry, parsed, packet, now, meta)
			if remove {
				s.removeEntry(key, entry, FlowCloseReset)
			}
			return handled
		}
	}
	judgement, judging := s.judging[key]
	if judging {
		if len(judgement.packets) < flowHoldPerFlow {
			s.hold(judgement, packet, parsed, meta)
		}
		return true
	}
	if parsed.protocol == uint8(header.TCPProtocolNumber) &&
		(parsed.tcpFlags&header.TCPFlagSyn == 0 || parsed.tcpFlags&header.TCPFlagAck != 0) {
		return false
	}
	judgement = &flowJudgement{key: key}
	if !s.hold(judgement, packet, parsed, meta) {
		return true
	}
	s.judging[key] = judgement
	var firstPacket []byte
	if parsed.protocol == uint8(header.UDPProtocolNumber) {
		firstPacket = bytes.Clone(header.UDP(parsed.transport).Payload())
	}
	go s.judge(judgement, firstPacket)
	return true
}

func (s *ForwardStage) hold(judgement *flowJudgement, packet []byte, parsed *forwardPacket, meta *ForwardFrameMeta) bool {
	if s.held >= flowHoldCapacity {
		return false
	}
	held := heldPacket{
		packet:          bytes.Clone(packet),
		transportLength: len(parsed.transport),
	}
	if meta != nil {
		held.meta = *meta
	}
	judgement.packets = append(judgement.packets, held)
	s.held++
	return true
}

func (s *ForwardStage) judge(judgement *flowJudgement, firstPacket []byte) {
	key := judgement.key
	verdict := s.dispatcher.handler.JudgeFlow(key.protocol, key.source, key.destination, firstPacket)
	if verdict.Action == ActionBypass && verdict.Port != nil {
		verdict.Action = ActionFlow
	}
	s.access.Lock()
	if s.judging[key] != judgement {
		s.access.Unlock()
		return
	}
	judgement.verdict = verdict
	s.judged = append(s.judged, judgement)
	s.access.Unlock()
	s.resume()
}

func (s *ForwardStage) ResumeJudged(process func(packet []byte, meta ForwardFrameMeta)) {
	var unhandled []heldPacket
	s.access.Lock()
	if s.dispatcher.closed.Load() {
		s.access.Unlock()
		return
	}
	now := s.dispatcher.now()
	for _, judgement := range s.judged {
		delete(s.judging, judgement.key)
		s.held -= len(judgement.packets)
		for index := range judgement.packets {
			held := &judgement.packets[index]
			parsed, _ := parseForwardPacket(held.packet)
			parsed.transport = parsed.transport[:held.transportLength]
			if !s.install(judgement.key, &parsed, held.packet, &held.meta, judgement.verdict, now) {
				unhandled = append(unhandled, *held)
			}
		}
	}
	clear(s.judged)
	s.judged = s.judged[:0]
	s.access.Unlock()
	for _, held := range unhandled {
		process(held.packet, held.meta)
	}
}

func (s *ForwardStage) teardownFlow(key flowKey, reason FlowCloseReason) {
	if s == nil || s.dispatcher.closed.Load() {
		return
	}
	s.access.Lock()
	entry, loaded := s.table[key]
	if loaded {
		s.removeEntry(key, entry, reason)
	}
	s.access.Unlock()
}

func (s *ForwardStage) handleHit(entry *flowEntry, packet *forwardPacket, raw []byte, now int64, meta *ForwardFrameMeta) (handled bool, remove bool) {
	switch entry.action {
	case ActionFlow:
		flow := entry.flow
		if flow.closed.Load() {
			tombstoneEntry(entry, now)
			return true, false
		}
		var flowFinished bool
		if packet.protocol == uint8(header.TCPProtocolNumber) {
			if packet.tcpFlags&header.TCPFlagRst != 0 {
				s.forwardToPort(flow, packet, raw, meta)
				flow.close(FlowCloseReset)
				tombstoneEntry(entry, now)
				return true, false
			}
			if packet.tcpFlags&header.TCPFlagFin != 0 {
				flow.finForward.Store(true)
			} else if flow.finForward.Load() && flow.finReverse.Load() {
				flowFinished = true
			}
		}
		entry.idle = s.dispatcher.flowIdle(flow)
		entry.deadline = now + int64(entry.idle)
		s.forwardToPort(flow, packet, raw, meta)
		if flowFinished {
			flow.report(FlowCloseFinished)
		}
		return true, false
	case ActionAccept:
		if packet.protocol == uint8(header.TCPProtocolNumber) {
			if packet.tcpFlags&header.TCPFlagSyn == 0 {
				return false, true
			}
		} else if packet.protocol == uint8(header.UDPProtocolNumber) {
			packet.verdict = entry.verdict
		}
		entry.deadline = now + int64(entry.idle)
		return false, false
	case ActionReject:
		entry.deadline = now + int64(entry.idle)
		s.stageReject(packet, raw, meta)
		return true, false
	default:
		entry.deadline = now + int64(entry.idle)
		return true, false
	}
}

func (s *ForwardStage) install(key flowKey, packet *forwardPacket, raw []byte, meta *ForwardFrameMeta, verdict FlowVerdict, now int64) bool {
	d := s.dispatcher
	existing, loaded := s.table[key]
	if loaded {
		if entryExpired(existing, now) {
			s.removeEntry(key, existing, FlowCloseTimeout)
		} else {
			handled, remove := s.handleHit(existing, packet, raw, now, meta)
			if remove {
				s.removeEntry(key, existing, FlowCloseReset)
			}
			return handled
		}
	}
	switch verdict.Action {
	case ActionFlow:
		if verdict.Port != nil {
			flow, result := s.createFlow(packet, verdict)
			if result == createFlowOK {
				entry := &flowEntry{action: ActionFlow, flow: flow, idle: d.flowIdle(flow)}
				entry.deadline = now + int64(entry.idle)
				s.insertEntry(key, entry, now)
				s.forwardToPort(flow, packet, raw, meta)
				return true
			}
			if result == createFlowExhausted {
				if now-s.exhaustedLogAt >= int64(exhaustedLogInterval) {
					s.exhaustedLogAt = now
					d.logger.Warn("port selector range exhausted, rejecting flow to ", packet.destination)
				}
				s.installSimple(key, ActionReject, packet.protocol, now)
				s.stageReject(packet, raw, meta)
				return true
			}
		}
		s.installAccept(key, packet, verdict, now)
		return false
	case ActionReject:
		s.installSimple(key, ActionReject, packet.protocol, now)
		s.stageReject(packet, raw, meta)
		return true
	case ActionDrop:
		s.installSimple(key, ActionDrop, packet.protocol, now)
		return true
	case ActionHijackDNS:
		if packet.protocol == uint8(header.UDPProtocolNumber) {
			d.hijackDNSPacket(packet)
			return true
		}
		s.installAccept(key, packet, verdict, now)
		return false
	default:
		s.installAccept(key, packet, verdict, now)
		return false
	}
}

func (s *ForwardStage) installSimple(key flowKey, action FlowAction, protocol uint8, now int64) *flowEntry {
	entry := &flowEntry{action: action, idle: s.dispatcher.idleTimeout(protocol, false)}
	entry.deadline = now + int64(entry.idle)
	s.insertEntry(key, entry, now)
	return entry
}

func (s *ForwardStage) installAccept(key flowKey, packet *forwardPacket, verdict FlowVerdict, now int64) {
	entry := s.installSimple(key, ActionAccept, packet.protocol, now)
	entry.verdict = verdict
	if packet.protocol == uint8(header.UDPProtocolNumber) {
		packet.verdict = verdict
	}
}

func (d *ForwardDispatcher) acceptedVerdict(protocol uint8, source netip.AddrPort, destination netip.AddrPort) (FlowVerdict, bool) {
	if d == nil || d.closed.Load() {
		return FlowVerdict{}, false
	}
	key := flowKey{protocol: protocol, source: source, destination: destination}
	now := d.now()
	d.stagesAccess.Lock()
	stages := d.stages
	d.stagesAccess.Unlock()
	for _, stage := range stages {
		stage.access.Lock()
		entry, loaded := stage.table[key]
		if loaded && entry.action == ActionAccept && !entryExpired(entry, now) {
			verdict := entry.verdict
			stage.access.Unlock()
			return verdict, true
		}
		stage.access.Unlock()
	}
	return FlowVerdict{}, false
}

func (d *ForwardDispatcher) idleTimeout(protocol uint8, established bool) time.Duration {
	switch protocol {
	case uint8(header.TCPProtocolNumber):
		if established {
			return tcpEstablishedTimeout
		}
		return tcpTransitoryTimeout
	case uint8(header.UDPProtocolNumber):
		return d.udpTimeout
	default:
		return d.icmpTimeout
	}
}

func (d *ForwardDispatcher) flowIdle(flow *forwardFlow) time.Duration {
	if flow.protocol == uint8(header.TCPProtocolNumber) && flow.finForward.Load() && flow.finReverse.Load() {
		return tcpClosingTimeout
	}
	if flow.udpTimeout > 0 {
		return flow.udpTimeout
	}
	established := flow.established.Load() && !flow.finForward.Load() && !flow.finReverse.Load()
	return d.idleTimeout(flow.protocol, established)
}

type createFlowResult uint8

const (
	createFlowOK createFlowResult = iota
	createFlowUnsupported
	createFlowExhausted
)

const exhaustedLogInterval = 5 * time.Second

func (s *ForwardStage) createFlow(packet *forwardPacket, verdict FlowVerdict) (*forwardFlow, createFlowResult) {
	var portAddress netip.Addr
	inet4Address, inet6Address := verdict.Port.PortAddresses()
	if packet.ipVersion == 6 {
		portAddress = inet6Address
	} else {
		portAddress = inet4Address
	}
	if !portAddress.IsValid() {
		return nil, createFlowUnsupported
	}
	effectiveMTU := verdict.Port.PortMTU()
	if s.dispatcher.mtu != 0 && (effectiveMTU == 0 || s.dispatcher.mtu < effectiveMTU) {
		effectiveMTU = s.dispatcher.mtu
	}
	if packet.ipVersion == 6 && effectiveMTU != 0 && effectiveMTU < header.IPv6MinimumMTU {
		return nil, createFlowUnsupported
	}
	isICMP := isICMPProtocol(packet.protocol)
	clientDestinationAddress := packet.destination.Addr()
	clientDestinationPort := packet.destination.Port()
	serverAddress := clientDestinationAddress
	serverPort := clientDestinationPort
	if verdict.Destination.Addr().IsValid() {
		serverAddress = verdict.Destination.Addr()
		if serverAddress.Is4() != clientDestinationAddress.Is4() {
			return nil, createFlowUnsupported
		}
	}
	if verdict.Destination.Port() != 0 && !isICMP {
		serverPort = verdict.Destination.Port()
	}
	nat := s.dispatcher.natFor(verdict.Port)
	if nat == nil {
		return nil, createFlowUnsupported
	}
	server := netip.AddrPortFrom(serverAddress, serverPort)
	natAddress := portAddress
	if nat.routed {
		natAddress = packet.source.Addr()
	}
	reservation, reserved := nat.reserve(s.dispatcher, packet.protocol, natAddress, packet.source, server)
	if !reserved {
		rangedPort, isRanged := portCapability[PortWithSelectorRange](verdict.Port)
		if isRanged && rangedPort.ExpandSelectorRanges(packet.protocol) {
			reservation, reserved = nat.reserve(s.dispatcher, packet.protocol, natAddress, packet.source, server)
		}
		if !reserved {
			return nil, createFlowExhausted
		}
	}
	var udpTimeout time.Duration
	if packet.protocol == uint8(header.UDPProtocolNumber) {
		udpTimeout = verdict.UDPTimeout
	}
	flow := &forwardFlow{
		nat:                      nat,
		owner:                    s,
		effectiveMTU:             effectiveMTU,
		protocol:                 packet.protocol,
		udpTimeout:               udpTimeout,
		clientAddress:            packet.source.Addr(),
		clientSelector:           packet.source.Port(),
		clientDestinationAddress: clientDestinationAddress,
		clientDestinationPort:    clientDestinationPort,
		serverAddress:            serverAddress,
		dnatAddress:              serverAddress != clientDestinationAddress,
		dnatPort:                 serverPort != clientDestinationPort && !isICMP,
	}
	flow.forwardRule = rewriteRule{rewriteSourcePort: true}
	if !nat.routed {
		flow.forwardRule.sourceAddress = addrToTCPIP(portAddress)
	}
	if flow.dnatAddress {
		flow.forwardRule.destinationAddress = addrToTCPIP(serverAddress)
	}
	if flow.dnatPort {
		flow.forwardRule.destinationPort = serverPort
		flow.forwardRule.rewriteDestinationPort = true
	}
	flow.reverseRule = rewriteRule{
		destinationAddress:     addrToTCPIP(flow.clientAddress),
		destinationPort:        flow.clientSelector,
		rewriteDestinationPort: true,
	}
	if flow.dnatAddress {
		flow.reverseRule.sourceAddress = addrToTCPIP(clientDestinationAddress)
	}
	if flow.dnatPort {
		flow.reverseRule.sourcePort = clientDestinationPort
		flow.reverseRule.rewriteSourcePort = true
	}
	if verdict.NewTracker != nil {
		flow.tracker = verdict.NewTracker()
		if flow.tracker != nil {
			flow.tracker.AttachFlow(flow)
		}
	}
	nat.publish(flow, reservation)
	return flow, createFlowOK
}

func (d *ForwardDispatcher) natFor(port Port) *portNAT {
	d.portsAccess.Lock()
	defer d.portsAccess.Unlock()
	returnPath, loaded := d.ports[port]
	if loaded {
		return returnPath.nat
	}
	returnPath = &portReturn{
		dispatcher: d,
		nat:        acquirePortNAT(port),
		fragments:  fragmentTable{entries: make(map[fragmentKey]*fragmentEntry)},
	}
	err := port.AttachReturn(returnPath)
	if err != nil {
		releasePortNAT(returnPath.nat)
		d.logger.Trace(E.Cause(err, "attach return path"))
		return nil
	}
	d.ports[port] = returnPath
	return returnPath.nat
}

type relayOutput interface {
	relayHeadroom() int
	relayReserve(count int, size int) ([][]byte, []int)
	relayPacket(packet []byte)
	relaySegment(segment []byte)
	relayTooBig(packet *forwardPacket, mtu uint32)
}

type stagePortOutput struct {
	stage    *ForwardStage
	nat      *portNAT
	borrowed bool
}

func (o *stagePortOutput) relayHeadroom() int {
	return 0
}

func (o *stagePortOutput) relayReserve(count int, size int) ([][]byte, []int) {
	return o.stage.reserveSegments(count, size)
}

func (o *stagePortOutput) relayPacket(packet []byte) {
	if o.borrowed {
		staged, _ := o.stage.reserveSegments(1, len(packet))
		copy(staged[0], packet)
		packet = staged[0]
	}
	o.stage.stagePort(o.nat, packet)
}

func (o *stagePortOutput) relaySegment(segment []byte) {
	o.stage.stagePort(o.nat, segment)
}

func (o *stagePortOutput) relayTooBig(packet *forwardPacket, mtu uint32) {
	o.stage.stageTooBig(packet, mtu)
}

type stageReturnOutput struct {
	stage     *ForwardStage
	writeback ForwardWriteback
}

func (o *stageReturnOutput) relayHeadroom() int {
	return o.writeback.ReturnHeadroom()
}

func (o *stageReturnOutput) relayReserve(count int, size int) ([][]byte, []int) {
	return o.stage.reserveSegments(count, size)
}

func (o *stageReturnOutput) relayPacket(packet []byte) {
	headroom := o.writeback.ReturnHeadroom()
	staged, _ := o.stage.reserveSegments(1, headroom+len(packet))
	copy(staged[0][headroom:], packet)
	o.stage.stageReturn(o.writeback, staged[0])
}

func (o *stageReturnOutput) relaySegment(segment []byte) {
	o.stage.stageReturn(o.writeback, segment)
}

func (o *stageReturnOutput) relayTooBig(packet *forwardPacket, mtu uint32) {
	o.stage.stageTooBig(packet, mtu)
}

func (s *ForwardStage) stageTooBig(packet *forwardPacket, mtu uint32) {
	s.writebackBatch = append(s.writebackBatch, buildICMPError(packet, ICMPErrorPacketTooBig, mtu, s.writeback.ReturnHeadroom()))
}

func (s *ForwardStage) forwardToPort(flow *forwardFlow, packet *forwardPacket, raw []byte, meta *ForwardFrameMeta) {
	s.portOutput.nat = flow.nat
	s.portOutput.borrowed = meta != nil
	relayFlow(&s.portOutput, s.dispatcher.logger, flow, &flow.forwardRule, packet, raw, meta, false)
}

func relayFlow(output relayOutput, log logger.Logger, flow *forwardFlow, rule *rewriteRule, packet *forwardPacket, raw []byte, meta *ForwardFrameMeta, reverse bool) {
	effectiveMTU := flow.effectiveMTU
	if meta != nil {
		meta.completeChecksum(raw)
		if meta.gsoSize != 0 && packet.protocol == uint8(header.TCPProtocolNumber) {
			headerLength := len(raw) - len(packet.transport) + int(header.TCP(packet.transport).DataOffset())
			offloadMTU := uint32(headerLength) + uint32(meta.gsoSize)
			if effectiveMTU == 0 || offloadMTU < effectiveMTU {
				effectiveMTU = offloadMTU
			}
		}
	}
	if effectiveMTU != 0 && uint32(len(raw)) > effectiveMTU {
		if packet.protocol == uint8(header.TCPProtocolNumber) {
			headerLength, segmentSize, segmentable := tcpSegmentSize(packet, raw, effectiveMTU)
			if segmentable {
				countRelay(flow, len(raw), reverse)
				rewritePacket(packet, rule, flow.effectiveMTU)
				err := segmentTCP(output, packet, raw, headerLength, segmentSize)
				if err != nil {
					log.Trace(E.Cause(err, "segment packet"))
				}
				return
			}
		} else if packet.ipVersion == 4 && header.IPv4(packet.network).Flags()&header.IPv4FlagDontFragment == 0 {
			countRelay(flow, len(raw), reverse)
			rewritePacket(packet, rule, flow.effectiveMTU)
			fragments, ok := fragmentIPv4Packet(header.IPv4(packet.network), flow.effectiveMTU, output.relayHeadroom())
			if ok {
				for _, fragment := range fragments {
					output.relaySegment(fragment)
				}
			}
			return
		}
		output.relayTooBig(packet, flow.effectiveMTU)
		return
	}
	countRelay(flow, len(raw), reverse)
	rewritePacket(packet, rule, flow.effectiveMTU)
	output.relayPacket(raw)
}

func countRelay(flow *forwardFlow, n int, reverse bool) {
	if flow.tracker == nil {
		return
	}
	if reverse {
		flow.tracker.CountReverse(n)
	} else {
		flow.tracker.CountForward(n)
	}
}

func rewritePacket(packet *forwardPacket, rule *rewriteRule, effectiveMTU uint32) {
	if rule == nil {
		return
	}
	if packet.fragment {
		if packet.hasFlow {
			applyRewrite(packet, rule)
		} else {
			applyRewriteAddresses(packet, rule)
		}
		return
	}
	if packet.isTCPSyn() {
		applyRewriteRaw(packet, rule)
		clampTCPMSS(packet, effectiveMTU)
		recomputeChecksums(packet)
	} else {
		applyRewrite(packet, rule)
	}
}

func (s *ForwardStage) stagePort(nat *portNAT, packet []byte) {
	for index := range s.ports {
		if s.ports[index].nat == nat {
			s.ports[index].packets = append(s.ports[index].packets, packet)
			return
		}
	}
	s.ports = append(s.ports, stagedPort{nat: nat, packets: [][]byte{packet}})
}

func (s *ForwardStage) stageReturn(writeback ForwardWriteback, packet []byte) {
	for index := range s.returns {
		if s.returns[index].writeback == writeback {
			s.returns[index].packets = append(s.returns[index].packets, packet)
			return
		}
	}
	s.returns = append(s.returns, returnBatch{writeback: writeback, packets: [][]byte{packet}})
}

func (s *ForwardStage) stageReject(packet *forwardPacket, raw []byte, meta *ForwardFrameMeta) {
	if meta != nil {
		meta.completeChecksum(raw)
	}
	reply, ok := buildReject(packet, s.writeback.ReturnHeadroom())
	if ok {
		s.writebackBatch = append(s.writebackBatch, reply)
	}
}

func (d *ForwardDispatcher) ResetNetwork() {
	if d == nil {
		return
	}
	d.resetGeneration.Add(1)
}

func (s *ForwardStage) Flush() {
	if s == nil {
		return
	}
	d := s.dispatcher
	if d.closed.Load() {
		return
	}
	now := d.now()
	generation := d.resetGeneration.Load()
	if s.resetSeen != generation {
		s.resetSeen = generation
		s.access.Lock()
		for key, entry := range s.table {
			s.removeEntry(key, entry, FlowCloseReset)
		}
		clear(s.judging)
		clear(s.judged)
		s.judged = s.judged[:0]
		s.held = 0
		s.access.Unlock()
	}
	for index := range s.ports {
		staged := &s.ports[index]
		if len(staged.packets) == 0 {
			continue
		}
		err := staged.nat.port.WritePackets(staged.packets)
		if err != nil {
			d.logger.Trace(E.Cause(err, "forward packets"))
		}
		clear(staged.packets)
		staged.packets = staged.packets[:0]
	}
	for index := range s.returns {
		staged := &s.returns[index]
		if len(staged.packets) == 0 {
			continue
		}
		err := staged.writeback.WriteReturnPackets(staged.packets)
		if err != nil {
			d.logger.Trace(E.Cause(err, "return packets"))
		}
		clear(staged.packets)
		staged.packets = staged.packets[:0]
	}
	if retain := max(s.segmentUsed, segmentRetainCount); len(s.segmentBuffers) > retain {
		clear(s.segmentBuffers[retain:])
		s.segmentBuffers = s.segmentBuffers[:retain]
		s.segmentSizes = s.segmentSizes[:retain]
	}
	s.segmentUsed = 0
	if len(s.writebackBatch) > 0 {
		err := s.writeback.WriteReturnPackets(s.writebackBatch)
		if err != nil {
			d.logger.Trace(E.Cause(err, "write back packets"))
		}
		clear(s.writebackBatch)
		s.writebackBatch = s.writebackBatch[:0]
	}
	if now-s.lastSweep >= int64(flowSweepInterval) {
		s.lastSweep = now
		s.access.Lock()
		s.sweep(now)
		s.access.Unlock()
	}
}

func (s *ForwardStage) sweepDue() (time.Duration, bool) {
	if s == nil {
		return 0, false
	}
	s.access.Lock()
	count := len(s.table)
	s.access.Unlock()
	if count == 0 {
		return 0, false
	}
	return max(time.Duration(s.lastSweep+int64(flowSweepInterval)-s.dispatcher.now()), 0), true
}

func entryExpired(entry *flowEntry, now int64) bool {
	if entry.flow != nil && entry.flow.aborted.Load() {
		return true
	}
	if now <= entry.deadline {
		return false
	}
	if entry.action == ActionFlow {
		lastReverse := entry.flow.lastReverse.Load()
		reverseDeadline := lastReverse + int64(entry.idle)
		if lastReverse != 0 && now <= reverseDeadline {
			entry.deadline = reverseDeadline
			return false
		}
	}
	return true
}

func tombstoneEntry(entry *flowEntry, now int64) {
	entry.action = ActionDrop
	entry.idle = flowTombstoneTimeout
	entry.deadline = now + int64(entry.idle)
}

func (s *ForwardStage) removeEntry(key flowKey, entry *flowEntry, reason FlowCloseReason) {
	delete(s.table, key)
	if entry.flow != nil {
		if reason == FlowCloseTimeout && entry.flow.finForward.Load() && entry.flow.finReverse.Load() {
			reason = FlowCloseFinished
		}
		entry.flow.close(reason)
		entry.flow.nat.release(entry.flow)
	}
}

func (s *ForwardStage) insertEntry(key flowKey, entry *flowEntry, now int64) {
	if len(s.table) >= max(flowTableCapacity/int(s.dispatcher.stageCount.Load()), flowStageCapacity) {
		s.evictEntries(now)
	}
	s.table[key] = entry
}

func (s *ForwardStage) evictEntries(now int64) {
	var (
		freed     int
		visited   int
		oldestKey flowKey
		oldest    *flowEntry
	)
	for key, entry := range s.table {
		if entryExpired(entry, now) {
			s.removeEntry(key, entry, FlowCloseTimeout)
			freed++
		} else if oldest == nil || entry.deadline < oldest.deadline {
			oldestKey = key
			oldest = entry
		}
		visited++
		if visited >= flowSweepLimit {
			break
		}
	}
	if freed == 0 && oldest != nil {
		s.removeEntry(oldestKey, oldest, FlowCloseReset)
	}
}

func (s *ForwardStage) sweep(now int64) {
	visited := 0
	for key, entry := range s.table {
		if entryExpired(entry, now) {
			s.removeEntry(key, entry, FlowCloseTimeout)
		} else if entry.action == ActionFlow && entry.flow.closed.Load() {
			tombstoneEntry(entry, now)
		}
		visited++
		if visited >= flowSweepLimit {
			break
		}
	}
}

func isICMPProtocol(protocol uint8) bool {
	return protocol == uint8(header.ICMPv4ProtocolNumber) || protocol == uint8(header.ICMPv6ProtocolNumber)
}

type portReturn struct {
	dispatcher *ForwardDispatcher
	nat        *portNAT
	fragments  fragmentTable
}

func (r *portReturn) ReturnHeadroom() int {
	return r.dispatcher.writeback.ReturnHeadroom()
}

type returnDecision uint8

const (
	returnPass returnDecision = iota
	returnWrite
	returnDrop
)

type returnBatch struct {
	writeback ForwardWriteback
	packets   [][]byte
}

type returnOutput struct {
	returnPath *portReturn
	headroom   int
	raw        []byte
	writeback  ForwardWriteback
	batches    []returnBatch
}

func (o *returnOutput) relayHeadroom() int {
	return o.headroom
}

func (o *returnOutput) relayReserve(count int, size int) ([][]byte, []int) {
	buffers := make([][]byte, count)
	for index := range buffers {
		buffers[index] = make([]byte, size)
	}
	return buffers, make([]int, count)
}

func (o *returnOutput) relayPacket(packet []byte) {
	o.relaySegment(o.raw)
}

func (o *returnOutput) relaySegment(segment []byte) {
	index := slices.IndexFunc(o.batches, func(batch returnBatch) bool { return batch.writeback == o.writeback })
	if index < 0 {
		o.batches = append(o.batches, returnBatch{writeback: o.writeback})
		index = len(o.batches) - 1
	}
	o.batches[index].packets = append(o.batches[index].packets, segment)
}

func (o *returnOutput) relayTooBig(packet *forwardPacket, mtu uint32) {
	err := o.returnPath.nat.port.WritePackets([][]byte{buildICMPError(packet, ICMPErrorPacketTooBig, mtu, 0)})
	if err != nil {
		o.returnPath.dispatcher.logger.Trace(E.Cause(err, "write packet too big"))
	}
}

func (o *returnOutput) flush() {
	for _, batch := range o.batches {
		err := batch.writeback.WriteReturnPackets(batch.packets)
		if err != nil {
			o.returnPath.dispatcher.logger.Trace(E.Cause(err, "write return packets"))
		}
	}
}

func (r *portReturn) ReturnPackets(packets [][]byte) [][]byte {
	if r.dispatcher.closed.Load() {
		return packets
	}
	output := returnOutput{returnPath: r, headroom: r.dispatcher.writeback.ReturnHeadroom()}
	now := r.dispatcher.now()
	unconsumed := packets[:0]
	for _, raw := range packets {
		if len(raw) < output.headroom+header.IPv4MinimumSize {
			unconsumed = append(unconsumed, raw)
			continue
		}
		parsed, ok := parseForwardPacket(raw[output.headroom:])
		if !ok {
			unconsumed = append(unconsumed, raw)
			continue
		}
		decision, flow, rule := r.classify(&parsed, now)
		switch decision {
		case returnWrite:
			output.raw = raw
			output.writeback = flow.owner.writeback
			relayFlow(&output, r.dispatcher.logger, flow, rule, &parsed, raw[output.headroom:], nil, true)
		case returnDrop:
		default:
			unconsumed = append(unconsumed, raw)
		}
	}
	output.flush()
	return unconsumed
}

func (r *portReturn) classify(parsed *forwardPacket, now int64) (returnDecision, *forwardFlow, *rewriteRule) {
	if parsed.fragment {
		return r.classifyFragment(parsed, now)
	}
	if !parsed.hasFlow {
		if parsed.isICMPError() {
			flow := r.returnICMPError(parsed)
			if flow != nil {
				return returnWrite, flow, nil
			}
		}
		return returnPass, nil, nil
	}
	flow, rule := r.matchReverse(parsed)
	if flow == nil {
		return returnPass, nil, nil
	}
	if flow.closed.Load() {
		return returnDrop, nil, nil
	}
	flow.observeReverse(parsed, now)
	return returnWrite, flow, rule
}

var errNoDispatcher = E.New("stack has no flow dispatcher")

func (d *ForwardDispatcher) attachReturn(returnPath Return) error {
	if d == nil {
		return errNoDispatcher
	}
	portReturnPath, isPortReturn := returnPath.(*portReturn)
	if !isPortReturn {
		return E.New("unsupported return path")
	}
	d.returnsAccess.Lock()
	defer d.returnsAccess.Unlock()
	current := *d.returns.Load()
	if slices.Contains(current, portReturnPath) {
		return nil
	}
	updated := append(slices.Clone(current), portReturnPath)
	d.returns.Store(&updated)
	return nil
}

func (d *ForwardDispatcher) detachReturn(returnPath Return) {
	if d == nil {
		return
	}
	portReturnPath, isPortReturn := returnPath.(*portReturn)
	if !isPortReturn {
		return
	}
	d.returnsAccess.Lock()
	defer d.returnsAccess.Unlock()
	current := *d.returns.Load()
	index := slices.Index(current, portReturnPath)
	if index < 0 {
		return
	}
	updated := slices.Delete(slices.Clone(current), index, index+1)
	d.returns.Store(&updated)
}

func (s *ForwardStage) ReturnParsed(packet []byte, meta ForwardFrameMeta, parsed *forwardPacket) bool {
	if s == nil || s.dispatcher.closed.Load() {
		return false
	}
	for _, returnPath := range *s.dispatcher.returns.Load() {
		decision, flow, rule := returnPath.classify(parsed, returnPath.dispatcher.now())
		switch decision {
		case returnWrite:
			s.returnOutput.writeback = flow.owner.writeback
			relayFlow(&s.returnOutput, s.dispatcher.logger, flow, rule, parsed, packet, &meta, true)
			return true
		case returnDrop:
			return true
		}
	}
	return false
}

func (r *portReturn) matchReverse(parsed *forwardPacket) (*forwardFlow, *rewriteRule) {
	key := parsed.flowKey()
	flow := r.nat.lookup(key)
	if flow != nil {
		if flow.owner.dispatcher != r.dispatcher {
			return nil, nil
		}
		return flow, &flow.reverseRule
	}
	if parsed.protocol != uint8(header.UDPProtocolNumber) {
		return nil, nil
	}
	mapping, anchor := r.nat.lookupMapping(key.destination.Addr(), key.destination.Port(), key.source.Addr())
	if mapping == nil || mapping.key.dispatcher != r.dispatcher {
		return nil, nil
	}
	return anchor, &mapping.reverseRule
}

func (r *portReturn) returnICMPError(parsed *forwardPacket) *forwardFlow {
	inner, ok := parsed.icmpErrorInner()
	if !ok {
		return nil
	}
	embedded, parsedInner := parseEmbedded(inner)
	if !parsedInner {
		return nil
	}
	flow := r.nat.lookup(embedded.flowKey().reversed())
	if flow == nil || flow.owner.dispatcher != r.dispatcher || flow.closed.Load() {
		return nil
	}
	rewriteEmbeddedSource(&embedded, addrToTCPIP(flow.clientAddress), flow.clientSelector, true)
	if flow.dnatAddress || flow.dnatPort {
		rewriteEmbeddedDestination(&embedded, addrToTCPIP(flow.clientDestinationAddress), flow.clientDestinationPort, flow.dnatPort)
	}
	sourceAddress, destinationAddress := parsed.addressSlices()
	copy(destinationAddress, flow.clientAddress.AsSlice())
	if parsed.source.Addr() == flow.serverAddress {
		copy(sourceAddress, flow.clientDestinationAddress.AsSlice())
	}
	recomputeChecksums(parsed)
	return flow
}