From 7c92d5e53e55cd57fb904b00d281ad7c2b5984a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Tue, 7 Jul 2026 21:19:56 +0800
Subject: [PATCH] Add UDPTimeout option for flow

---
 flow.go          |  6 +++++-
 flow_dispatch.go | 14 +++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/flow.go b/flow.go
index db2cd1d..4fbf801 100644
--- a/flow.go
+++ b/flow.go
@@ -1,11 +1,15 @@
 package tun
 
-import "net/netip"
+import (
+	"net/netip"
+	"time"
+)
 
 type FlowVerdict struct {
 	Action      FlowAction
 	Port        Port
 	Destination netip.AddrPort
+	UDPTimeout  time.Duration
 	NewTracker  func() FlowTracker
 }
 
diff --git a/flow_dispatch.go b/flow_dispatch.go
index e35a521..152a5bb 100644
--- a/flow_dispatch.go
+++ b/flow_dispatch.go
@@ -1,6 +1,7 @@
 package tun
 
 import (
+	"maps"
 	"net/netip"
 	"sync/atomic"
 	"time"
@@ -47,6 +48,7 @@ type forwardFlow struct {
 	reverseRule  rewriteRule
 	effectiveMTU uint32
 	protocol     uint8
+	udpTimeout   time.Duration
 	tracker      FlowTracker
 
 	clientAddress            netip.Addr
@@ -293,6 +295,9 @@ func (d *ForwardDispatcher) flowIdle(flow *forwardFlow) time.Duration {
 	if flow.protocol == uint8(header.TCPProtocolNumber) && flow.finForward && flow.finReverse.Load() {
 		return tcpClosingTimeout
 	}
+	if flow.udpTimeout > 0 {
+		return flow.udpTimeout
+	}
 	established := flow.established.Load() && !flow.finForward && !flow.finReverse.Load()
 	return d.idleTimeout(flow.protocol, established)
 }
@@ -331,11 +336,16 @@ func (d *ForwardDispatcher) createFlow(packet *forwardPacket, verdict FlowVerdic
 	if !allocated {
 		return nil, false
 	}
+	var udpTimeout time.Duration
+	if packet.protocol == uint8(header.UDPProtocolNumber) {
+		udpTimeout = verdict.UDPTimeout
+	}
 	flow := &forwardFlow{
 		nat:                      nat,
 		reverseKey:               reverseKey,
 		effectiveMTU:             effectiveMTU,
 		protocol:                 packet.protocol,
+		udpTimeout:               udpTimeout,
 		clientAddress:            packet.source.Addr(),
 		clientSelector:           packet.source.Port(),
 		clientDestinationAddress: clientDestinationAddress,
@@ -399,9 +409,7 @@ func (d *ForwardDispatcher) natFor(port Port) *portNAT {
 	d.natList.Store(&natList)
 	revMap := make(map[netip.Addr]*portNAT)
 	if currentRev := d.revNAT.Load(); currentRev != nil {
-		for addr, existing := range *currentRev {
-			revMap[addr] = existing
-		}
+		maps.Copy(revMap, *currentRev)
 	}
 	v4Address, v6Address := port.PortAddresses()
 	if v4Address.IsValid() {