Add closer wrapper

This commit is contained in:
世界 2022-08-17 20:46:08 +08:00
parent 81504dc2c4
commit ce854cda85
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
3 changed files with 76 additions and 0 deletions

View file

@ -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 ./...

18
common/closer.go Normal file
View file

@ -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}
}

View file

@ -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
}