Update udpnat usage

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/fsnotify/fsnotify v1.6.0
 	github.com/sagernet/go-tun2socks v1.16.12-0.20220818015926-16cb67876a61
 	github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97
-	github.com/sagernet/sing v0.2.4-0.20230418025125-f196b4303e31
+	github.com/sagernet/sing v0.2.4-0.20230419150837-2b3a62786474
 	golang.org/x/net v0.8.0
 	golang.org/x/sys v0.7.0
 	gvisor.dev/gvisor v0.0.0-20220901235040-6ca97ef2ce1c

go.sum

@@ -7,8 +7,8 @@ github.com/sagernet/go-tun2socks v1.16.12-0.20220818015926-16cb67876a61/go.mod h
 github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97 h1:iL5gZI3uFp0X6EslacyapiRz7LLSJyr4RajF/BhMVyE=
 github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97/go.mod h1:xLnfdiJbSp8rNqYEdIW/6eDO4mVoogml14Bh2hSiFpM=
 github.com/sagernet/sing v0.0.0-20220817130738-ce854cda8522/go.mod h1:QVsS5L/ZA2Q5UhQwLrn0Trw+msNd/NPGEhBKR/ioWiY=
-github.com/sagernet/sing v0.2.4-0.20230418025125-f196b4303e31 h1:qgq8jeY/rbnY9NwYXByO//AP0ByIxnsKUxQx1tOB3W0=
-github.com/sagernet/sing v0.2.4-0.20230418025125-f196b4303e31/go.mod h1:Ta8nHnDLAwqySzKhGoKk4ZIB+vJ3GTKj7UPrWYvM+4w=
+github.com/sagernet/sing v0.2.4-0.20230419150837-2b3a62786474 h1:eSYMHrZvHo9hTKSAPTFaGsiafUss7FPhTDanXsNrfwE=
+github.com/sagernet/sing v0.2.4-0.20230419150837-2b3a62786474/go.mod h1:Ta8nHnDLAwqySzKhGoKk4ZIB+vJ3GTKj7UPrWYvM+4w=
 github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 h1:gga7acRE695APm9hlsSMoOoE65U4/TcqNj90mc69Rlg=
 github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74/go.mod h1:DD4vA1DwXk04H54A1oHXtwZmA0grkVMdPxx/VGLCah0=
 golang.org/x/net v0.8.0 h1:Zrh2ngAOFYneWTAIAPethzeaQLuHwhuBkuV6ZiRnUaQ=

gvisor.go

@@ -8,6 +8,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/bufio"
 	"github.com/sagernet/sing/common/canceler"
 	E "github.com/sagernet/sing/common/exceptions"
@@ -42,6 +43,7 @@ type GVisor struct {
 	stack                  *stack.Stack
 	endpoint               stack.LinkEndpoint
 	routeMapping           *RouteMapping
+	udpForwarder           *UDPForwarder
 }
 
 type GVisorTun interface {
@@ -68,7 +70,7 @@ func NewGVisor(
 		logger:                 options.Logger,
 	}
 	if gStack.router != nil {
-		gStack.routeMapping = NewRouteMapping(options.Context, options.UDPTimeout)
+		gStack.routeMapping = NewRouteMapping(options.UDPTimeout)
 	}
 	return gStack, nil
 }
@@ -256,7 +258,8 @@ func (t *GVisor) Start() error {
 			return udpForwarder.HandlePacket(id, buffer)
 		})
 	} else {
-		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, NewUDPForwarder(t.ctx, ipStack, t.handler, t.udpTimeout).HandlePacket)
+		t.udpForwarder = NewUDPForwarder(t.ctx, ipStack, t.handler, t.udpTimeout)
+		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, t.udpForwarder.HandlePacket)
 	}
 
 	t.stack = ipStack
@@ -270,5 +273,8 @@ func (t *GVisor) Close() error {
 	for _, endpoint := range t.stack.CleanupEndpoints() {
 		endpoint.Abort()
 	}
-	return nil
+	return common.Close(
+		common.PtrOrNil(t.routeMapping),
+		common.PtrOrNil(t.udpForwarder),
+	)
 }

gvisor_udp.go

@@ -29,10 +29,14 @@ func NewUDPForwarder(ctx context.Context, stack *stack.Stack, handler Handler, u
 	return &UDPForwarder{
 		ctx:    ctx,
 		stack:  stack,
-		udpNat: udpnat.New[netip.AddrPort](ctx, udpTimeout, handler),
+		udpNat: udpnat.New[netip.AddrPort](udpTimeout, handler),
 	}
 }
 
+func (f *UDPForwarder) Close() error {
+	return f.udpNat.Close()
+}
+
 func (f *UDPForwarder) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
 	var upstreamMetadata M.Metadata
 	upstreamMetadata.Source = M.SocksaddrFrom(M.AddrFromIP(net.IP(id.RemoteAddress)), id.RemotePort)

lwip.go

@@ -35,7 +35,7 @@ func NewLWIP(
 		tunMtu:  options.MTU,
 		handler: options.Handler,
 		stack:   lwip.NewLWIPStack(),
-		udpNat:  udpnat.New[netip.AddrPort](options.Context, options.UDPTimeout, options.Handler),
+		udpNat:  udpnat.New[netip.AddrPort](options.UDPTimeout, options.Handler),
 	}, nil
 }
 
@@ -96,7 +96,10 @@ func (l *LWIP) Close() error {
 	lwip.RegisterOutputFn(func(bytes []byte) (int, error) {
 		return 0, os.ErrClosed
 	})
-	return l.stack.Close()
+	return common.Close(
+		l.stack,
+		common.PtrOrNil(l.udpNat),
+	)
 }
 
 func (l *LWIP) Handle(conn net.Conn) error {

route_mapping.go

@@ -2,16 +2,19 @@ package tun
 
 import (
 	"context"
+	"net"
 
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/cache"
 )
 
 type RouteMapping struct {
 	status *cache.LruCache[RouteSession, RouteAction]
+	cancel common.ContextCancelCauseFunc
 }
 
-func NewRouteMapping(ctx context.Context, maxAge int64) *RouteMapping {
+func NewRouteMapping(maxAge int64) *RouteMapping {
+	ctx, cancel := common.ContextWithCancelCause(context.Background())
 	return &RouteMapping{
 		status: cache.New(
 			cache.WithContext[RouteSession, RouteAction](ctx),
@@ -21,6 +24,7 @@ func NewRouteMapping(ctx context.Context, maxAge int64) *RouteMapping {
 				common.Close(conn)
 			}),
 		),
+		cancel: cancel,
 	}
 }
 
@@ -33,3 +37,8 @@ func (m *RouteMapping) Lookup(session RouteSession, constructor func() RouteActi
 	}
 	return action
 }
+
+func (m *RouteMapping) Close() error {
+	m.cancel(net.ErrClosed)
+	return nil
+}

system.go

@@ -65,7 +65,7 @@ func NewSystem(options StackOptions) (Stack, error) {
 		underPlatform: options.UnderPlatform,
 	}
 	if stack.router != nil {
-		stack.routeMapping = NewRouteMapping(options.Context, options.UDPTimeout)
+		stack.routeMapping = NewRouteMapping(options.UDPTimeout)
 	}
 	if len(options.Inet4Address) > 0 {
 		if options.Inet4Address[0].Bits() == 32 {
@@ -91,6 +91,8 @@ func (s *System) Close() error {
 	return common.Close(
 		s.tcpListener,
 		s.tcpListener6,
+		s.udpNat,
+		common.PtrOrNil(s.routeMapping),
 	)
 }
 
@@ -118,7 +120,7 @@ func (s *System) Start() error {
 		go s.acceptLoop(tcpListener)
 	}
 	s.tcpNat = NewNat(s.ctx, time.Second*time.Duration(s.udpTimeout))
-	s.udpNat = udpnat.New[netip.AddrPort](s.ctx, s.udpTimeout, s.handler)
+	s.udpNat = udpnat.New[netip.AddrPort](s.udpTimeout, s.handler)
 	go s.tunLoop()
 	return nil
 }