From 16a35527b7eff8a78fb9c0a68867e7e87e90057e Mon Sep 17 00:00:00 2001
From: yaling888 <73897884+yaling888@users.noreply.github.com>
Date: Sun, 24 Apr 2022 02:19:23 +0800
Subject: [PATCH] Chore: increase nattable capacity

---
 listener/tun/ipstack/system/mars/nat/table.go | 45 +++++++++++++------
 listener/tun/ipstack/system/mars/nat/tcp.go   | 10 +++++
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/listener/tun/ipstack/system/mars/nat/table.go b/listener/tun/ipstack/system/mars/nat/table.go
index e0d86ccc..9c1b32cd 100644
--- a/listener/tun/ipstack/system/mars/nat/table.go
+++ b/listener/tun/ipstack/system/mars/nat/table.go
@@ -1,13 +1,15 @@
 package nat
 
 import (
-	"container/list"
 	"net/netip"
+	"sync"
+
+	"github.com/Dreamacro/clash/common/generics/list"
 )
 
 const (
 	portBegin  = 30000
-	portLength = 4096
+	portLength = 10240
 )
 
 var zeroTuple = tuple{}
@@ -23,9 +25,10 @@ type binding struct {
 }
 
 type table struct {
-	tuples    map[tuple]*list.Element
-	ports     [portLength]*list.Element
-	available *list.List
+	mu        sync.Mutex
+	tuples    map[tuple]*list.Element[*binding]
+	ports     [portLength]*list.Element[*binding]
+	available *list.List[*binding]
 }
 
 func (t *table) tupleOf(port uint16) tuple {
@@ -36,28 +39,31 @@ func (t *table) tupleOf(port uint16) tuple {
 
 	elm := t.ports[offset]
 
-	t.available.MoveToFront(elm)
-
-	return elm.Value.(*binding).tuple
+	return elm.Value.tuple
 }
 
 func (t *table) portOf(tuple tuple) uint16 {
+	t.mu.Lock()
 	elm := t.tuples[tuple]
+	t.mu.Unlock()
 	if elm == nil {
 		return 0
 	}
 
 	t.available.MoveToFront(elm)
 
-	return portBegin + elm.Value.(*binding).offset
+	return portBegin + elm.Value.offset
 }
 
 func (t *table) newConn(tuple tuple) uint16 {
 	elm := t.available.Back()
-	b := elm.Value.(*binding)
+	b := elm.Value
 
+	t.mu.Lock()
 	delete(t.tuples, b.tuple)
 	t.tuples[tuple] = elm
+	t.mu.Unlock()
+
 	b.tuple = tuple
 
 	t.available.MoveToFront(elm)
@@ -65,11 +71,24 @@ func (t *table) newConn(tuple tuple) uint16 {
 	return portBegin + b.offset
 }
 
+func (t *table) delete(tup tuple) {
+	t.mu.Lock()
+	elm := t.tuples[tup]
+	if elm == nil {
+		t.mu.Unlock()
+		return
+	}
+	delete(t.tuples, tup)
+	t.mu.Unlock()
+
+	t.available.MoveToBack(elm)
+}
+
 func newTable() *table {
 	result := &table{
-		tuples:    make(map[tuple]*list.Element, portLength),
-		ports:     [portLength]*list.Element{},
-		available: list.New(),
+		tuples:    make(map[tuple]*list.Element[*binding], portLength),
+		ports:     [portLength]*list.Element[*binding]{},
+		available: list.New[*binding](),
 	}
 
 	for idx := range result.ports {
diff --git a/listener/tun/ipstack/system/mars/nat/tcp.go b/listener/tun/ipstack/system/mars/nat/tcp.go
index cc0abe7d..48ad3e43 100644
--- a/listener/tun/ipstack/system/mars/nat/tcp.go
+++ b/listener/tun/ipstack/system/mars/nat/tcp.go
@@ -16,6 +16,8 @@ type conn struct {
 	net.Conn
 
 	tuple tuple
+
+	close func(tuple tuple)
 }
 
 func (t *TCP) Accept() (net.Conn, error) {
@@ -37,6 +39,9 @@ func (t *TCP) Accept() (net.Conn, error) {
 	return &conn{
 		Conn:  c,
 		tuple: tup,
+		close: func(tuple tuple) {
+			t.table.delete(tuple)
+		},
 	}, nil
 }
 
@@ -52,6 +57,11 @@ func (t *TCP) SetDeadline(time time.Time) error {
 	return t.listener.SetDeadline(time)
 }
 
+func (c *conn) Close() error {
+	c.close(c.tuple)
+	return c.Conn.Close()
+}
+
 func (c *conn) LocalAddr() net.Addr {
 	return &net.TCPAddr{
 		IP:   c.tuple.SourceAddr.Addr().AsSlice(),