From d5dd3b3bb14aa53ecdf24b0694b45793fe0d63d3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sat, 30 Jul 2022 14:12:11 +0800
Subject: [PATCH] Fix leak endpoint

---
 go.mod    |  2 +-
 go.sum    |  4 ++--
 gvisor.go | 30 ++++++++++++++++++------------
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/go.mod b/go.mod
index eb929b1..f144de0 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/sagernet/sing-tun
 go 1.18
 
 require (
-	github.com/sagernet/sing v0.0.0-20220729120910-4376f188c512
+	github.com/sagernet/sing v0.0.0-20220730061053-a21e329a2698
 	github.com/vishvananda/netlink v1.1.0
 	golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10
 	gvisor.dev/gvisor v0.0.0-20220711011657-cecae2f4234d
diff --git a/go.sum b/go.sum
index 9767470..dbd25f4 100644
--- a/go.sum
+++ b/go.sum
@@ -1,7 +1,7 @@
 github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
 github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
-github.com/sagernet/sing v0.0.0-20220729120910-4376f188c512 h1:dCWDE55LpZu//W02FccNbGObZFlv1N2NS0yUdf2i4Mc=
-github.com/sagernet/sing v0.0.0-20220729120910-4376f188c512/go.mod h1:GbtQfZSpmtD3cXeD1qX2LCMwY8dH+bnnInDTqd92IsM=
+github.com/sagernet/sing v0.0.0-20220730061053-a21e329a2698 h1:wjoF4/FOwze8cN2/EvQyyuq1tzXjxNViPIoqQ7CNIb8=
+github.com/sagernet/sing v0.0.0-20220730061053-a21e329a2698/go.mod h1:GbtQfZSpmtD3cXeD1qX2LCMwY8dH+bnnInDTqd92IsM=
 github.com/vishvananda/netlink v1.1.0 h1:1iyaYNBLmP6L0220aDnYQpo1QEV4t4hJ+xEEhhJH8j0=
 github.com/vishvananda/netlink v1.1.0/go.mod h1:cTgwzPIzzgDAYoQrMm0EdrjRUBkTqKYppBueQtXaqoE=
 github.com/vishvananda/netns v0.0.0-20191106174202-0a2b9b5464df/go.mod h1:JP3t17pCcGlemwknint6hfoeCVQrEMVwxRLRjXpq+BU=
diff --git a/gvisor.go b/gvisor.go
index 5962a60..cf4b675 100644
--- a/gvisor.go
+++ b/gvisor.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"time"
 
-	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/bufio"
 	E "github.com/sagernet/sing/common/exceptions"
 	M "github.com/sagernet/sing/common/metadata"
@@ -96,7 +95,16 @@ func (t *GVisorTun) Start() error {
 
 	ipStack.SetTransportProtocolHandler(tcp.ProtocolNumber, tcp.NewForwarder(ipStack, 0, 1024, func(r *tcp.ForwarderRequest) {
 		var wq waiter.Queue
+		handshakeCtx, cancel := context.WithCancel(context.Background())
+		go func() {
+			select {
+			case <-t.ctx.Done():
+				wq.Notify(wq.Events())
+			case <-handshakeCtx.Done():
+			}
+		}()
 		endpoint, err := r.CreateEndpoint(&wq)
+		cancel()
 		if err != nil {
 			r.Complete(true)
 			return
@@ -118,10 +126,8 @@ func (t *GVisorTun) Start() error {
 			var metadata M.Metadata
 			metadata.Source = M.SocksaddrFromNet(lAddr)
 			metadata.Destination = M.SocksaddrFromNet(rAddr)
-			hErr := t.handler.NewConnection(t.ctx, &gTCPConn{tcpConn}, metadata)
-			if hErr != nil {
-				endpoint.Abort()
-			}
+			t.handler.NewConnection(t.ctx, &gTCPConn{tcpConn}, metadata)
+			endpoint.Abort()
 		}()
 	}).HandlePacket)
 
@@ -143,10 +149,8 @@ func (t *GVisorTun) Start() error {
 				var metadata M.Metadata
 				metadata.Source = M.SocksaddrFromNet(lAddr)
 				metadata.Destination = M.SocksaddrFromNet(rAddr)
-				hErr := t.handler.NewPacketConnection(ContextWithNeedTimeout(t.ctx, true), bufio.NewPacketConn(&bufio.UnbindPacketConn{ExtendedConn: bufio.NewExtendedConn(&gUDPConn{udpConn}), Addr: M.SocksaddrFromNet(rAddr)}), metadata)
-				if hErr != nil {
-					endpoint.Abort()
-				}
+				t.handler.NewPacketConnection(ContextWithNeedTimeout(t.ctx, true), bufio.NewPacketConn(&bufio.UnbindPacketConn{ExtendedConn: bufio.NewExtendedConn(&gUDPConn{udpConn}), Addr: M.SocksaddrFromNet(rAddr)}), metadata)
+				endpoint.Abort()
 			}()
 		}).HandlePacket)
 	} else {
@@ -158,7 +162,9 @@ func (t *GVisorTun) Start() error {
 }
 
 func (t *GVisorTun) Close() error {
-	return common.Close(
-		common.PtrOrNil(t.stack),
-	)
+	t.stack.Close()
+	for _, endpoint := range t.stack.CleanupEndpoints() {
+		endpoint.Abort()
+	}
+	return nil
 }