bufio: Refactor copy

This commit is contained in:
世界
2026-02-03 02:51:05 +08:00
parent c98fb6615a
commit 400aa437fd
3 changed files with 51 additions and 106 deletions

2
go.mod
View File

@@ -33,7 +33,7 @@ require (
github.com/sagernet/gomobile v0.1.11 github.com/sagernet/gomobile v0.1.11
github.com/sagernet/gvisor v0.0.0-20250811.0-sing-box-mod.1 github.com/sagernet/gvisor v0.0.0-20250811.0-sing-box-mod.1
github.com/sagernet/quic-go v0.59.0-sing-box-mod.2 github.com/sagernet/quic-go v0.59.0-sing-box-mod.2
github.com/sagernet/sing v0.8.0-beta.15 github.com/sagernet/sing v0.8.0-beta.15.0.20260202162209-7c477e13f41e
github.com/sagernet/sing-mux v0.3.4 github.com/sagernet/sing-mux v0.3.4
github.com/sagernet/sing-quic v0.6.0-beta.11 github.com/sagernet/sing-quic v0.6.0-beta.11
github.com/sagernet/sing-shadowsocks v0.2.8 github.com/sagernet/sing-shadowsocks v0.2.8

4
go.sum
View File

@@ -210,8 +210,8 @@ github.com/sagernet/nftables v0.3.0-beta.4 h1:kbULlAwAC3jvdGAC1P5Fa3GSxVwQJibNen
github.com/sagernet/nftables v0.3.0-beta.4/go.mod h1:OQXAjvjNGGFxaTgVCSTRIhYB5/llyVDeapVoENYBDS8= github.com/sagernet/nftables v0.3.0-beta.4/go.mod h1:OQXAjvjNGGFxaTgVCSTRIhYB5/llyVDeapVoENYBDS8=
github.com/sagernet/quic-go v0.59.0-sing-box-mod.2 h1:hJUL+HtxEOjxsa0CsucbBVqI/AMS4k52NwNU637zmdw= github.com/sagernet/quic-go v0.59.0-sing-box-mod.2 h1:hJUL+HtxEOjxsa0CsucbBVqI/AMS4k52NwNU637zmdw=
github.com/sagernet/quic-go v0.59.0-sing-box-mod.2/go.mod h1:OqILvS182CyOol5zNNo6bguvOGgXzV459+chpRaUC+4= github.com/sagernet/quic-go v0.59.0-sing-box-mod.2/go.mod h1:OqILvS182CyOol5zNNo6bguvOGgXzV459+chpRaUC+4=
github.com/sagernet/sing v0.8.0-beta.15 h1:lP6XnzeQvVBfuTkByo5YnG4Oy/AVkDC2ZljghSfHzKQ= github.com/sagernet/sing v0.8.0-beta.15.0.20260202162209-7c477e13f41e h1:H8izpW6d9l8Ub5UFSV/Q2WCehss2KAlmnDiABa4BHp0=
github.com/sagernet/sing v0.8.0-beta.15/go.mod h1:ARkL0gM13/Iv5VCZmci/NuoOlePoIsW0m7BWfln/Hak= github.com/sagernet/sing v0.8.0-beta.15.0.20260202162209-7c477e13f41e/go.mod h1:ARkL0gM13/Iv5VCZmci/NuoOlePoIsW0m7BWfln/Hak=
github.com/sagernet/sing-mux v0.3.4 h1:ZQplKl8MNXutjzbMVtWvWG31fohhgOfCuUZR4dVQ8+s= github.com/sagernet/sing-mux v0.3.4 h1:ZQplKl8MNXutjzbMVtWvWG31fohhgOfCuUZR4dVQ8+s=
github.com/sagernet/sing-mux v0.3.4/go.mod h1:QvlKMyNBNrQoyX4x+gq028uPbLM2XeRpWtDsWBJbFSk= github.com/sagernet/sing-mux v0.3.4/go.mod h1:QvlKMyNBNrQoyX4x+gq028uPbLM2XeRpWtDsWBJbFSk=
github.com/sagernet/sing-quic v0.6.0-beta.11 h1:eUusxITKKRedhWC2ScUYFUvD96h/QfbKLaS3N6/7in4= github.com/sagernet/sing-quic v0.6.0-beta.11 h1:eUusxITKKRedhWC2ScUYFUvD96h/QfbKLaS3N6/7in4=

View File

@@ -2,7 +2,6 @@ package route
import ( import (
"context" "context"
"errors"
"io" "io"
"net" "net"
"net/netip" "net/netip"
@@ -102,8 +101,12 @@ func (m *ConnectionManager) NewConnection(ctx context.Context, this N.Dialer, co
m.connections.Remove(element) m.connections.Remove(element)
}) })
var done atomic.Bool var done atomic.Bool
m.preConnectionCopy(ctx, conn, remoteConn, false, &done, onClose) if m.kickWriteHandshake(ctx, conn, remoteConn, false, &done, onClose) {
m.preConnectionCopy(ctx, remoteConn, conn, true, &done, onClose) return
}
if m.kickWriteHandshake(ctx, remoteConn, conn, true, &done, onClose) {
return
}
go m.connectionCopy(ctx, conn, remoteConn, false, &done, onClose) go m.connectionCopy(ctx, conn, remoteConn, false, &done, onClose)
go m.connectionCopy(ctx, remoteConn, conn, true, &done, onClose) go m.connectionCopy(ctx, remoteConn, conn, true, &done, onClose)
} }
@@ -226,75 +229,8 @@ func (m *ConnectionManager) NewPacketConnection(ctx context.Context, this N.Dial
go m.packetConnectionCopy(ctx, destination, conn, true, &done, onClose) go m.packetConnectionCopy(ctx, destination, conn, true, &done, onClose)
} }
func (m *ConnectionManager) preConnectionCopy(ctx context.Context, source net.Conn, destination net.Conn, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) {
readHandshake := N.NeedHandshakeForRead(source)
writeHandshake := N.NeedHandshakeForWrite(destination)
if readHandshake || writeHandshake {
var err error
for {
err = m.connectionCopyEarlyWrite(source, destination, readHandshake, writeHandshake)
if err == nil && N.NeedHandshakeForRead(source) {
continue
} else if E.IsMulti(err, os.ErrInvalid, context.DeadlineExceeded, io.EOF) {
err = nil
}
break
}
if err != nil {
if done.Swap(true) {
onClose(err)
}
common.Close(source, destination)
if !direction {
m.logger.ErrorContext(ctx, "connection upload handshake: ", err)
} else {
m.logger.ErrorContext(ctx, "connection download handshake: ", err)
}
return
}
}
}
func (m *ConnectionManager) connectionCopy(ctx context.Context, source net.Conn, destination net.Conn, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) { func (m *ConnectionManager) connectionCopy(ctx context.Context, source net.Conn, destination net.Conn, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) {
var ( _, err := bufio.CopyWithIncreateBuffer(destination, source, bufio.DefaultIncreaseBufferAfter, bufio.DefaultBatchSize)
sourceReader io.Reader = source
destinationWriter io.Writer = destination
)
var readCounters, writeCounters []N.CountFunc
for {
sourceReader, readCounters = N.UnwrapCountReader(sourceReader, readCounters)
destinationWriter, writeCounters = N.UnwrapCountWriter(destinationWriter, writeCounters)
if cachedSrc, isCached := sourceReader.(N.CachedReader); isCached {
cachedBuffer := cachedSrc.ReadCached()
if cachedBuffer != nil {
dataLen := cachedBuffer.Len()
_, err := destination.Write(cachedBuffer.Bytes())
cachedBuffer.Release()
if err != nil {
if done.Swap(true) {
onClose(err)
}
common.Close(source, destination)
if !direction {
m.logger.ErrorContext(ctx, "connection upload payload: ", err)
} else {
m.logger.ErrorContext(ctx, "connection download payload: ", err)
}
return
}
for _, counter := range readCounters {
counter(int64(dataLen))
}
for _, counter := range writeCounters {
counter(int64(dataLen))
}
}
continue
}
break
}
_, err := bufio.CopyWithCounters(destinationWriter, sourceReader, source, readCounters, writeCounters, bufio.DefaultIncreaseBufferAfter, bufio.DefaultBatchSize)
if err != nil { if err != nil {
common.Close(source, destination) common.Close(source, destination)
} else if duplexDst, isDuplex := destination.(N.WriteCloser); isDuplex { } else if duplexDst, isDuplex := destination.(N.WriteCloser); isDuplex {
@@ -328,45 +264,54 @@ func (m *ConnectionManager) connectionCopy(ctx context.Context, source net.Conn,
} }
} }
func (m *ConnectionManager) connectionCopyEarlyWrite(source net.Conn, destination io.Writer, readHandshake bool, writeHandshake bool) error { func (m *ConnectionManager) kickWriteHandshake(ctx context.Context, source net.Conn, destination net.Conn, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) bool {
payload := buf.NewPacket() if !N.NeedHandshakeForWrite(destination) {
defer payload.Release() return false
err := source.SetReadDeadline(time.Now().Add(C.ReadPayloadTimeout))
if err != nil {
if err == os.ErrInvalid {
if writeHandshake {
return common.Error(destination.Write(nil))
}
}
return err
} }
var ( var (
isTimeout bool cachedBuffer *buf.Buffer
isEOF bool wrotePayload bool
) )
_, err = payload.ReadOnceFrom(source) sourceReader, readCounters := N.UnwrapCountReader(source, nil)
if err != nil { destinationWriter, writeCounters := N.UnwrapCountWriter(destination, nil)
if E.IsTimeout(err) { if cachedReader, ok := sourceReader.(N.CachedReader); ok {
isTimeout = true cachedBuffer = cachedReader.ReadCached()
} else if errors.Is(err, io.EOF) { }
isEOF = true var err error
} else { if cachedBuffer != nil {
return E.Cause(err, "read payload") wrotePayload = true
dataLen := cachedBuffer.Len()
_, err = destinationWriter.Write(cachedBuffer.Bytes())
cachedBuffer.Release()
if err == nil {
for _, counter := range readCounters {
counter(int64(dataLen))
}
for _, counter := range writeCounters {
counter(int64(dataLen))
}
} }
} else {
_ = destination.SetWriteDeadline(time.Now().Add(C.ReadPayloadTimeout))
_, err = destinationWriter.Write(nil)
_ = destination.SetWriteDeadline(time.Time{})
} }
_ = source.SetReadDeadline(time.Time{}) if err == nil {
if !payload.IsEmpty() || writeHandshake { return false
_, err = destination.Write(payload.Bytes())
if err != nil {
return E.Cause(err, "write payload")
}
} }
if isTimeout { if !wrotePayload && (E.IsMulti(err, os.ErrInvalid, context.DeadlineExceeded, io.EOF) || E.IsTimeout(err)) {
return context.DeadlineExceeded return false
} else if isEOF {
return io.EOF
} }
return nil if !done.Swap(true) {
onClose(err)
}
common.Close(source, destination)
if !direction {
m.logger.ErrorContext(ctx, "connection upload handshake: ", err)
} else {
m.logger.ErrorContext(ctx, "connection download handshake: ", err)
}
return true
} }
func (m *ConnectionManager) packetConnectionCopy(ctx context.Context, source N.PacketReader, destination N.PacketWriter, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) { func (m *ConnectionManager) packetConnectionCopy(ctx context.Context, source N.PacketReader, destination N.PacketWriter, direction bool, done *atomic.Bool, onClose N.CloseHandlerFunc) {