diff --git a/Makefile b/Makefile index 6f7c6f1..f7a8a53 100644 --- a/Makefile +++ b/Makefile @@ -9,6 +9,7 @@ fmt_install: lint: GOOS=linux golangci-lint run ./... + GOOS=android golangci-lint run ./... GOOS=windows golangci-lint run ./... GOOS=darwin golangci-lint run ./... GOOS=freebsd golangci-lint run ./... diff --git a/common/closer.go b/common/closer.go new file mode 100644 index 0000000..cc0b6d1 --- /dev/null +++ b/common/closer.go @@ -0,0 +1,18 @@ +package common + +import "io" + +type closeWrapper struct { + closer func() error +} + +func (w *closeWrapper) Close() error { + return w.closer() +} + +func Closer(closer func() error) io.Closer { + if closer == nil { + return nil + } + return &closeWrapper{closer} +} diff --git a/common/conntrack/tracker.go b/common/conntrack/tracker.go index b6c7662..429447f 100644 --- a/common/conntrack/tracker.go +++ b/common/conntrack/tracker.go @@ -2,9 +2,12 @@ package conntrack import ( "io" + "net" "sync" "github.com/sagernet/sing/common" + "github.com/sagernet/sing/common/bufio" + N "github.com/sagernet/sing/common/network" "github.com/sagernet/sing/common/x/list" ) @@ -20,6 +23,16 @@ func (m *Tracker) Track(conn io.Closer) *Registration { return &Registration{m, element} } +func (m *Tracker) TrackConn(conn net.Conn) net.Conn { + registration := m.Track(conn) + return &trackConn{conn, registration} +} + +func (m *Tracker) TrackPacketConn(conn net.PacketConn) N.NetPacketConn { + registration := m.Track(conn) + return &trackPacketConn{bufio.NewPacketConn(conn), registration} +} + func (m *Tracker) Reset() { m.access.Lock() defer m.access.Unlock() @@ -39,3 +52,47 @@ func (t *Registration) Leave() { defer t.manager.access.Unlock() t.manager.connections.Remove(t.element) } + +type trackConn struct { + net.Conn + registration *Registration +} + +func (t *trackConn) Close() error { + t.registration.Leave() + return t.Conn.Close() +} + +func (t *trackConn) WriteTo(w io.Writer) (n int64, err error) { + return bufio.Copy(w, t.Conn) +} + +func (t *trackConn) ReadFrom(r io.Reader) (n int64, err error) { + return bufio.Copy(t.Conn, r) +} + +func (t *trackConn) Upstream() any { + return t.Conn +} + +func (t *trackConn) ReaderReplaceable() bool { + return true +} + +func (t *trackConn) WriterReplaceable() bool { + return true +} + +type trackPacketConn struct { + N.NetPacketConn + registration *Registration +} + +func (t *trackPacketConn) Close() error { + t.registration.Leave() + return t.NetPacketConn.Close() +} + +func (t *trackPacketConn) Upstream() any { + return t.NetPacketConn +}