Improve direct copy

This commit is contained in:
世界
2023-04-24 19:01:10 +08:00
parent bc32c78d03
commit c287731df9
7 changed files with 114 additions and 197 deletions

View File

@@ -7,9 +7,9 @@ import (
"time"
"github.com/sagernet/sing-box/adapter"
"github.com/sagernet/sing-box/experimental/trackerconn"
"github.com/sagernet/sing/common"
"github.com/sagernet/sing/common/atomic"
"github.com/sagernet/sing/common/bufio"
N "github.com/sagernet/sing/common/network"
"github.com/gofrs/uuid/v5"
@@ -115,13 +115,13 @@ func NewTCPTracker(conn net.Conn, manager *Manager, metadata Metadata, router ad
download := new(atomic.Int64)
t := &tcpTracker{
ExtendedConn: trackerconn.NewHook(conn, func(n int64) {
ExtendedConn: bufio.NewCounterConn(conn, []N.CountFunc{func(n int64) {
upload.Add(n)
manager.PushUploaded(n)
}, func(n int64) {
}}, []N.CountFunc{func(n int64) {
download.Add(n)
manager.PushDownloaded(n)
}),
}}),
manager: manager,
trackerInfo: &trackerInfo{
UUID: uuid,
@@ -202,13 +202,13 @@ func NewUDPTracker(conn N.PacketConn, manager *Manager, metadata Metadata, route
download := new(atomic.Int64)
ut := &udpTracker{
PacketConn: trackerconn.NewHookPacket(conn, func(n int64) {
PacketConn: bufio.NewCounterPacketConn(conn, []N.CountFunc{func(n int64) {
upload.Add(n)
manager.PushUploaded(n)
}, func(n int64) {
}}, []N.CountFunc{func(n int64) {
download.Add(n)
manager.PushDownloaded(n)
}),
}}),
manager: manager,
trackerInfo: &trackerInfo{
UUID: uuid,

View File

@@ -1,108 +0,0 @@
package trackerconn
import (
"net"
"github.com/sagernet/sing/common/atomic"
"github.com/sagernet/sing/common/buf"
"github.com/sagernet/sing/common/bufio"
N "github.com/sagernet/sing/common/network"
)
func New(conn net.Conn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64) *Conn {
return &Conn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
}
func NewHook(conn net.Conn, readCounter func(n int64), writeCounter func(n int64)) *HookConn {
return &HookConn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
}
type Conn struct {
N.ExtendedConn
readCounter []*atomic.Int64
writeCounter []*atomic.Int64
}
func (c *Conn) Read(p []byte) (n int, err error) {
n, err = c.ExtendedConn.Read(p)
for _, counter := range c.readCounter {
counter.Add(int64(n))
}
return n, err
}
func (c *Conn) ReadBuffer(buffer *buf.Buffer) error {
err := c.ExtendedConn.ReadBuffer(buffer)
if err != nil {
return err
}
for _, counter := range c.readCounter {
counter.Add(int64(buffer.Len()))
}
return nil
}
func (c *Conn) Write(p []byte) (n int, err error) {
n, err = c.ExtendedConn.Write(p)
for _, counter := range c.writeCounter {
counter.Add(int64(n))
}
return n, err
}
func (c *Conn) WriteBuffer(buffer *buf.Buffer) error {
dataLen := int64(buffer.Len())
err := c.ExtendedConn.WriteBuffer(buffer)
if err != nil {
return err
}
for _, counter := range c.writeCounter {
counter.Add(dataLen)
}
return nil
}
func (c *Conn) Upstream() any {
return c.ExtendedConn
}
type HookConn struct {
N.ExtendedConn
readCounter func(n int64)
writeCounter func(n int64)
}
func (c *HookConn) Read(p []byte) (n int, err error) {
n, err = c.ExtendedConn.Read(p)
c.readCounter(int64(n))
return n, err
}
func (c *HookConn) ReadBuffer(buffer *buf.Buffer) error {
err := c.ExtendedConn.ReadBuffer(buffer)
if err != nil {
return err
}
c.readCounter(int64(buffer.Len()))
return nil
}
func (c *HookConn) Write(p []byte) (n int, err error) {
n, err = c.ExtendedConn.Write(p)
c.writeCounter(int64(n))
return n, err
}
func (c *HookConn) WriteBuffer(buffer *buf.Buffer) error {
dataLen := int64(buffer.Len())
err := c.ExtendedConn.WriteBuffer(buffer)
if err != nil {
return err
}
c.writeCounter(dataLen)
return nil
}
func (c *HookConn) Upstream() any {
return c.ExtendedConn
}

View File

@@ -1,76 +0,0 @@
package trackerconn
import (
"github.com/sagernet/sing/common/atomic"
"github.com/sagernet/sing/common/buf"
M "github.com/sagernet/sing/common/metadata"
N "github.com/sagernet/sing/common/network"
)
func NewPacket(conn N.PacketConn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64) *PacketConn {
return &PacketConn{conn, readCounter, writeCounter}
}
func NewHookPacket(conn N.PacketConn, readCounter func(n int64), writeCounter func(n int64)) *HookPacketConn {
return &HookPacketConn{conn, readCounter, writeCounter}
}
type PacketConn struct {
N.PacketConn
readCounter []*atomic.Int64
writeCounter []*atomic.Int64
}
func (c *PacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
destination, err = c.PacketConn.ReadPacket(buffer)
if err == nil {
for _, counter := range c.readCounter {
counter.Add(int64(buffer.Len()))
}
}
return
}
func (c *PacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
dataLen := int64(buffer.Len())
err := c.PacketConn.WritePacket(buffer, destination)
if err != nil {
return err
}
for _, counter := range c.writeCounter {
counter.Add(dataLen)
}
return nil
}
func (c *PacketConn) Upstream() any {
return c.PacketConn
}
type HookPacketConn struct {
N.PacketConn
readCounter func(n int64)
writeCounter func(n int64)
}
func (c *HookPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
destination, err = c.PacketConn.ReadPacket(buffer)
if err == nil {
c.readCounter(int64(buffer.Len()))
}
return
}
func (c *HookPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
dataLen := int64(buffer.Len())
err := c.PacketConn.WritePacket(buffer, destination)
if err != nil {
return err
}
c.writeCounter(dataLen)
return nil
}
func (c *HookPacketConn) Upstream() any {
return c.PacketConn
}

View File

@@ -10,9 +10,9 @@ import (
"time"
"github.com/sagernet/sing-box/adapter"
"github.com/sagernet/sing-box/experimental/trackerconn"
"github.com/sagernet/sing-box/option"
"github.com/sagernet/sing/common/atomic"
"github.com/sagernet/sing/common/bufio"
E "github.com/sagernet/sing/common/exceptions"
N "github.com/sagernet/sing/common/network"
)
@@ -83,7 +83,7 @@ func (s *StatsService) RoutedConnection(inbound string, outbound string, user st
writeCounter = append(writeCounter, s.loadOrCreateCounter("user>>>"+user+">>>traffic>>>downlink"))
}
s.access.Unlock()
return trackerconn.New(conn, readCounter, writeCounter)
return bufio.NewInt64CounterConn(conn, readCounter, writeCounter)
}
func (s *StatsService) RoutedPacketConnection(inbound string, outbound string, user string, conn N.PacketConn) N.PacketConn {
@@ -109,7 +109,7 @@ func (s *StatsService) RoutedPacketConnection(inbound string, outbound string, u
writeCounter = append(writeCounter, s.loadOrCreateCounter("user>>>"+user+">>>traffic>>>downlink"))
}
s.access.Unlock()
return trackerconn.NewPacket(conn, readCounter, writeCounter)
return bufio.NewInt64CounterPacketConn(conn, readCounter, writeCounter)
}
func (s *StatsService) GetStats(ctx context.Context, request *GetStatsRequest) (*GetStatsResponse, error) {