Improve conntrack

common/conntrack/conn.go

@@ -0,0 +1,54 @@
+package conntrack
+
+import (
+	"io"
+	"net"
+
+	"github.com/sagernet/sing/common/x/list"
+)
+
+type Conn struct {
+	net.Conn
+	element *list.Element[io.Closer]
+}
+
+func NewConn(conn net.Conn) (net.Conn, error) {
+	connAccess.Lock()
+	element := openConnection.PushBack(conn)
+	connAccess.Unlock()
+	if KillerEnabled {
+		err := KillerCheck()
+		if err != nil {
+			conn.Close()
+			return nil, err
+		}
+	}
+	return &Conn{
+		Conn:    conn,
+		element: element,
+	}, nil
+}
+
+func (c *Conn) Close() error {
+	if c.element.Value != nil {
+		connAccess.Lock()
+		if c.element.Value != nil {
+			openConnection.Remove(c.element)
+			c.element.Value = nil
+		}
+		connAccess.Unlock()
+	}
+	return c.Conn.Close()
+}
+
+func (c *Conn) Upstream() any {
+	return c.Conn
+}
+
+func (c *Conn) ReaderReplaceable() bool {
+	return true
+}
+
+func (c *Conn) WriterReplaceable() bool {
+	return true
+}

common/conntrack/killer.go

@@ -0,0 +1,35 @@
+package conntrack
+
+import (
+	runtimeDebug "runtime/debug"
+	"time"
+
+	E "github.com/sagernet/sing/common/exceptions"
+	"github.com/sagernet/sing/common/memory"
+)
+
+var (
+	KillerEnabled   bool
+	MemoryLimit     uint64
+	killerLastCheck time.Time
+)
+
+func KillerCheck() error {
+	if !KillerEnabled {
+		return nil
+	}
+	nowTime := time.Now()
+	if nowTime.Sub(killerLastCheck) < 3*time.Second {
+		return nil
+	}
+	killerLastCheck = nowTime
+	if memory.Total() > MemoryLimit {
+		Close()
+		go func() {
+			time.Sleep(time.Second)
+			runtimeDebug.FreeOSMemory()
+		}()
+		return E.New("out of memory")
+	}
+	return nil
+}

common/conntrack/packet_conn.go

@@ -0,0 +1,55 @@
+package conntrack
+
+import (
+	"io"
+	"net"
+
+	"github.com/sagernet/sing/common/bufio"
+	"github.com/sagernet/sing/common/x/list"
+)
+
+type PacketConn struct {
+	net.PacketConn
+	element *list.Element[io.Closer]
+}
+
+func NewPacketConn(conn net.PacketConn) (net.PacketConn, error) {
+	connAccess.Lock()
+	element := openConnection.PushBack(conn)
+	connAccess.Unlock()
+	if KillerEnabled {
+		err := KillerCheck()
+		if err != nil {
+			conn.Close()
+			return nil, err
+		}
+	}
+	return &PacketConn{
+		PacketConn: conn,
+		element:    element,
+	}, nil
+}
+
+func (c *PacketConn) Close() error {
+	if c.element.Value != nil {
+		connAccess.Lock()
+		if c.element.Value != nil {
+			openConnection.Remove(c.element)
+			c.element.Value = nil
+		}
+		connAccess.Unlock()
+	}
+	return c.PacketConn.Close()
+}
+
+func (c *PacketConn) Upstream() any {
+	return bufio.NewPacketConn(c.PacketConn)
+}
+
+func (c *PacketConn) ReaderReplaceable() bool {
+	return true
+}
+
+func (c *PacketConn) WriterReplaceable() bool {
+	return true
+}

common/conntrack/track.go

@@ -0,0 +1,47 @@
+package conntrack
+
+import (
+	"io"
+	"sync"
+
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/x/list"
+)
+
+var (
+	connAccess     sync.RWMutex
+	openConnection list.List[io.Closer]
+)
+
+func Count() int {
+	if !Enabled {
+		return 0
+	}
+	return openConnection.Len()
+}
+
+func List() []io.Closer {
+	if !Enabled {
+		return nil
+	}
+	connAccess.RLock()
+	defer connAccess.RUnlock()
+	connList := make([]io.Closer, 0, openConnection.Len())
+	for element := openConnection.Front(); element != nil; element = element.Next() {
+		connList = append(connList, element.Value)
+	}
+	return connList
+}
+
+func Close() {
+	if !Enabled {
+		return
+	}
+	connAccess.Lock()
+	defer connAccess.Unlock()
+	for element := openConnection.Front(); element != nil; element = element.Next() {
+		common.Close(element.Value)
+		element.Value = nil
+	}
+	openConnection.Init()
+}

common/conntrack/track_disable.go

@@ -0,0 +1,5 @@
+//go:build !with_conntrack
+
+package conntrack
+
+const Enabled = false

common/conntrack/track_enable.go

@@ -0,0 +1,5 @@
+//go:build with_conntrack
+
+package conntrack
+
+const Enabled = true