Fix buffer usage

This commit is contained in:
世界 2022-08-01 20:48:08 +08:00
parent 1bb95f9661
commit 842fd713ff
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
10 changed files with 15 additions and 61 deletions

View file

@ -532,6 +532,7 @@ type ExtendedWriterWrapper struct {
}
func (w *ExtendedWriterWrapper) WriteBuffer(buffer *buf.Buffer) error {
defer buffer.Release()
return common.Error(w.Write(buffer.Bytes()))
}

View file

@ -22,7 +22,6 @@ func ReadFrom(reader N.PacketReader, buffer *buf.Buffer) (n int, addr net.Addr,
}
func Write(writer N.ExtendedWriter, buffer *buf.Buffer) (n int, err error) {
defer buffer.Release()
dataLen := buffer.Len()
err = writer.WriteBuffer(buffer)
if err == nil {
@ -32,7 +31,6 @@ func Write(writer N.ExtendedWriter, buffer *buf.Buffer) (n int, err error) {
}
func WriteTo(writer N.PacketWriter, buffer *buf.Buffer, addr net.Addr) (n int, err error) {
defer buffer.Release()
dataLen := buffer.Len()
err = writer.WritePacket(buffer, M.SocksaddrFromNet(addr))
if err == nil {

View file

@ -1,50 +0,0 @@
package bufio
import (
"io"
"github.com/sagernet/sing/common"
"github.com/sagernet/sing/common/buf"
N "github.com/sagernet/sing/common/network"
)
type LimitedWriter struct {
upstream N.ExtendedWriter
maxChunkLength int
}
func NewLimitedWriter(writer io.Writer, maxChunkLength int) *LimitedWriter {
return &LimitedWriter{
upstream: NewExtendedWriter(writer),
maxChunkLength: maxChunkLength,
}
}
func (w *LimitedWriter) Write(p []byte) (n int, err error) {
for pLen := len(p); pLen > 0; {
var data []byte
if pLen > w.maxChunkLength {
data = p[:w.maxChunkLength]
p = p[w.maxChunkLength:]
pLen -= w.maxChunkLength
} else {
data = p
pLen = 0
}
var writeN int
writeN, err = w.upstream.Write(data)
if err != nil {
return
}
n += writeN
}
return
}
func (w *LimitedWriter) WriteBuffer(buffer *buf.Buffer) error {
if buffer.Len() <= w.maxChunkLength {
return w.upstream.WriteBuffer(buffer)
}
defer buffer.Release()
return common.Error(w.Write(buffer.Bytes()))
}

View file

@ -22,7 +22,7 @@ func CopyOnce(dst io.Writer, src io.Reader) (n int64, err error) {
func CopyExtendedOnce(dst N.ExtendedWriter, src N.ExtendedReader) (n int64, err error) {
var buffer *buf.Buffer
if _, unsafe := common.Cast[N.ThreadUnsafeWriter](dst); unsafe {
if N.IsUnsafeWriter(dst) {
buffer = buf.New()
} else {
_buffer := buf.StackNew()

View file

@ -22,6 +22,11 @@ type HeadroomWriter interface {
Headroom() int
}
func IsUnsafeWriter(writer any) bool {
_, isUnsafe := common.Cast[ThreadUnsafeWriter](writer)
return isUnsafe
}
func CalculateHeadroom(writer any) int {
var headroom int
if headroomWriter, needHeadroom := writer.(HeadroomWriter); needHeadroom {

View file

@ -44,6 +44,7 @@ func (c *ServerConn) RemoteAddr() net.Addr {
return pipeAddr{}
}
//warn:unsafe
func (c *ServerConn) loopInput() {
_buffer := buf.StackNew()
defer common.KeepAlive(_buffer)
@ -79,6 +80,7 @@ func (c *ServerConn) loopInput() {
c.Close()
}
//warn:unsafe
func (c *ServerConn) loopOutput() {
_buffer := buf.StackNew()
defer common.KeepAlive(_buffer)

2
go.mod
View file

@ -2,4 +2,4 @@ module github.com/sagernet/sing
go 1.18
require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8
require golang.org/x/sys v0.0.0-20220731174439-a90be440212d

4
go.sum
View file

@ -1,2 +1,2 @@
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220731174439-a90be440212d h1:Sv5ogFZatcgIMMtBSTTAgMYsicp25MXBubjXNDKwm80=
golang.org/x/sys v0.0.0-20220731174439-a90be440212d/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

View file

@ -63,10 +63,7 @@ func HandleConnection(ctx context.Context, conn net.Conn, reader *std_bufio.Read
var requestConn net.Conn
if reader.Buffered() > 0 {
_buffer := buf.StackNewSize(reader.Buffered())
defer common.KeepAlive(_buffer)
buffer := common.Dup(_buffer)
defer buffer.Release()
buffer := buf.NewSize(reader.Buffered())
_, err = buffer.ReadFullFrom(reader, reader.Buffered())
if err != nil {
return err

View file

@ -43,6 +43,7 @@ func (c *AssociatePacketConn) RemoteAddr() net.Addr {
return c.addr
}
//warn:unsafe
func (c *AssociatePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
buffer := buf.With(p)
n, _, err = bufio.ReadFrom(c.PacketConn, buffer)
@ -59,6 +60,7 @@ func (c *AssociatePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err erro
return
}
//warn:unsafe
func (c *AssociatePacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
destination := M.SocksaddrFromNet(addr)
_buffer := buf.StackNewSize(3 + M.SocksaddrSerializer.AddrPortLen(destination) + len(p))
@ -98,7 +100,6 @@ func (c *AssociatePacketConn) ReadPacket(buffer *buf.Buffer) (M.Socksaddr, error
}
func (c *AssociatePacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
defer buffer.Release()
header := buf.With(buffer.ExtendHeader(3 + M.SocksaddrSerializer.AddrPortLen(destination)))
common.Must(header.WriteZeroN(3))
common.Must(M.SocksaddrSerializer.WriteAddrPort(header, destination))