Add NDIS inbound

common/conntrack/conn.go

@@ -1,54 +0,0 @@
-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/context.go

@@ -0,0 +1,14 @@
+package conntrack
+
+import (
+	"context"
+
+	"github.com/sagernet/sing/service"
+)
+
+func ContextWithDefaultTracker(ctx context.Context, killerEnabled bool, memoryLimit uint64) context.Context {
+	if service.FromContext[Tracker](ctx) != nil {
+		return ctx
+	}
+	return service.ContextWith[Tracker](ctx, NewDefaultTracker(killerEnabled, memoryLimit))
+}

common/conntrack/default.go

@@ -0,0 +1,245 @@
+package conntrack
+
+import (
+	"net"
+	"net/netip"
+	runtimeDebug "runtime/debug"
+	"sync"
+	"time"
+
+	"github.com/sagernet/sing/common"
+	E "github.com/sagernet/sing/common/exceptions"
+	"github.com/sagernet/sing/common/memory"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
+	"github.com/sagernet/sing/common/x/list"
+)
+
+var _ Tracker = (*DefaultTracker)(nil)
+
+type DefaultTracker struct {
+	connAccess  sync.RWMutex
+	connList    list.List[net.Conn]
+	connAddress map[netip.AddrPort]netip.AddrPort
+
+	packetConnAccess  sync.RWMutex
+	packetConnList    list.List[AbstractPacketConn]
+	packetConnAddress map[netip.AddrPort]bool
+
+	pendingAccess sync.RWMutex
+	pendingList   list.List[netip.AddrPort]
+
+	killerEnabled   bool
+	memoryLimit     uint64
+	killerLastCheck time.Time
+}
+
+func NewDefaultTracker(killerEnabled bool, memoryLimit uint64) *DefaultTracker {
+	return &DefaultTracker{
+		connAddress:       make(map[netip.AddrPort]netip.AddrPort),
+		packetConnAddress: make(map[netip.AddrPort]bool),
+		killerEnabled:     killerEnabled,
+		memoryLimit:       memoryLimit,
+	}
+}
+
+func (t *DefaultTracker) NewConn(conn net.Conn) (net.Conn, error) {
+	err := t.KillerCheck()
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	t.connAccess.Lock()
+	element := t.connList.PushBack(conn)
+	t.connAddress[M.AddrPortFromNet(conn.LocalAddr())] = M.AddrPortFromNet(conn.RemoteAddr())
+	t.connAccess.Unlock()
+	return &Conn{
+		Conn: conn,
+		closeFunc: common.OnceFunc(func() {
+			t.removeConn(element)
+		}),
+	}, nil
+}
+
+func (t *DefaultTracker) NewConnEx(conn net.Conn) (N.CloseHandlerFunc, error) {
+	err := t.KillerCheck()
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	t.connAccess.Lock()
+	element := t.connList.PushBack(conn)
+	t.connAddress[M.AddrPortFromNet(conn.LocalAddr())] = M.AddrPortFromNet(conn.RemoteAddr())
+	t.connAccess.Unlock()
+	return N.OnceClose(func(it error) {
+		t.removeConn(element)
+	}), nil
+}
+
+func (t *DefaultTracker) NewPacketConn(conn net.PacketConn) (net.PacketConn, error) {
+	err := t.KillerCheck()
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	t.packetConnAccess.Lock()
+	element := t.packetConnList.PushBack(conn)
+	t.packetConnAddress[M.AddrPortFromNet(conn.LocalAddr())] = true
+	t.packetConnAccess.Unlock()
+	return &PacketConn{
+		PacketConn: conn,
+		closeFunc: common.OnceFunc(func() {
+			t.removePacketConn(element)
+		}),
+	}, nil
+}
+
+func (t *DefaultTracker) NewPacketConnEx(conn AbstractPacketConn) (N.CloseHandlerFunc, error) {
+	err := t.KillerCheck()
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	t.packetConnAccess.Lock()
+	element := t.packetConnList.PushBack(conn)
+	t.packetConnAddress[M.AddrPortFromNet(conn.LocalAddr())] = true
+	t.packetConnAccess.Unlock()
+	return N.OnceClose(func(it error) {
+		t.removePacketConn(element)
+	}), nil
+}
+
+func (t *DefaultTracker) CheckConn(source netip.AddrPort, destination netip.AddrPort) bool {
+	t.connAccess.RLock()
+	defer t.connAccess.RUnlock()
+	return t.connAddress[source] == destination
+}
+
+func (t *DefaultTracker) CheckPacketConn(source netip.AddrPort) bool {
+	t.packetConnAccess.RLock()
+	defer t.packetConnAccess.RUnlock()
+	return t.packetConnAddress[source]
+}
+
+func (t *DefaultTracker) AddPendingDestination(destination netip.AddrPort) func() {
+	t.pendingAccess.Lock()
+	defer t.pendingAccess.Unlock()
+	element := t.pendingList.PushBack(destination)
+	return func() {
+		t.pendingAccess.Lock()
+		defer t.pendingAccess.Unlock()
+		t.pendingList.Remove(element)
+	}
+}
+
+func (t *DefaultTracker) CheckDestination(destination netip.AddrPort) bool {
+	t.pendingAccess.RLock()
+	defer t.pendingAccess.RUnlock()
+	for element := t.pendingList.Front(); element != nil; element = element.Next() {
+		if element.Value == destination {
+			return true
+		}
+	}
+	return false
+}
+
+func (t *DefaultTracker) KillerCheck() error {
+	if !t.killerEnabled {
+		return nil
+	}
+	nowTime := time.Now()
+	if nowTime.Sub(t.killerLastCheck) < 3*time.Second {
+		return nil
+	}
+	t.killerLastCheck = nowTime
+	if memory.Total() > t.memoryLimit {
+		t.Close()
+		go func() {
+			time.Sleep(time.Second)
+			runtimeDebug.FreeOSMemory()
+		}()
+		return E.New("out of memory")
+	}
+	return nil
+}
+
+func (t *DefaultTracker) Count() int {
+	t.connAccess.RLock()
+	defer t.connAccess.RUnlock()
+	t.packetConnAccess.RLock()
+	defer t.packetConnAccess.RUnlock()
+	return t.connList.Len() + t.packetConnList.Len()
+}
+
+func (t *DefaultTracker) Close() {
+	t.connAccess.Lock()
+	for element := t.connList.Front(); element != nil; element = element.Next() {
+		element.Value.Close()
+	}
+	t.connList.Init()
+	t.connAccess.Unlock()
+	t.packetConnAccess.Lock()
+	for element := t.packetConnList.Front(); element != nil; element = element.Next() {
+		element.Value.Close()
+	}
+	t.packetConnList.Init()
+	t.packetConnAccess.Unlock()
+}
+
+func (t *DefaultTracker) removeConn(element *list.Element[net.Conn]) {
+	t.connAccess.Lock()
+	defer t.connAccess.Unlock()
+	delete(t.connAddress, M.AddrPortFromNet(element.Value.LocalAddr()))
+	t.connList.Remove(element)
+}
+
+func (t *DefaultTracker) removePacketConn(element *list.Element[AbstractPacketConn]) {
+	t.packetConnAccess.Lock()
+	defer t.packetConnAccess.Unlock()
+	delete(t.packetConnAddress, M.AddrPortFromNet(element.Value.LocalAddr()))
+	t.packetConnList.Remove(element)
+}
+
+type Conn struct {
+	net.Conn
+	closeFunc func()
+}
+
+func (c *Conn) Close() error {
+	c.closeFunc()
+	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
+}
+
+type PacketConn struct {
+	net.PacketConn
+	closeFunc func()
+}
+
+func (c *PacketConn) Close() error {
+	c.closeFunc()
+	return c.PacketConn.Close()
+}
+
+func (c *PacketConn) Upstream() any {
+	return c.PacketConn
+}
+
+func (c *PacketConn) ReaderReplaceable() bool {
+	return true
+}
+
+func (c *PacketConn) WriterReplaceable() bool {
+	return true
+}

common/conntrack/killer.go

@@ -1,35 +0,0 @@
-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

@@ -1,55 +0,0 @@
-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

@@ -1,47 +0,0 @@
-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

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

common/conntrack/track_enable.go

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

common/conntrack/tracker.go

@@ -0,0 +1,32 @@
+package conntrack
+
+import (
+	"net"
+	"net/netip"
+	"time"
+
+	N "github.com/sagernet/sing/common/network"
+)
+
+// TODO: add to N
+type AbstractPacketConn interface {
+	Close() error
+	LocalAddr() net.Addr
+	SetDeadline(t time.Time) error
+	SetReadDeadline(t time.Time) error
+	SetWriteDeadline(t time.Time) error
+}
+
+type Tracker interface {
+	NewConn(conn net.Conn) (net.Conn, error)
+	NewPacketConn(conn net.PacketConn) (net.PacketConn, error)
+	NewConnEx(conn net.Conn) (N.CloseHandlerFunc, error)
+	NewPacketConnEx(conn AbstractPacketConn) (N.CloseHandlerFunc, error)
+	CheckConn(source netip.AddrPort, destination netip.AddrPort) bool
+	CheckPacketConn(source netip.AddrPort) bool
+	AddPendingDestination(destination netip.AddrPort) func()
+	CheckDestination(destination netip.AddrPort) bool
+	KillerCheck() error
+	Count() int
+	Close()
+}