SagerNet/sing
1
0
Fork 0
mirror of https://github.com/SagerNet/sing.git synced 2025-04-05 04:47:40 +03:00
Code Issues Projects Releases Packages Wiki Activity Actions

Refine buffered reader

Browse source
This commit is contained in:
世界 2022-06-12 11:30:07 +08:00
parent 857116b2cc
commit 82cc7d29b8
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
2 changed files with 34 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
common/bufio/buffer.go
View file

@ -29,6 +29,10 @@ func (r *BufferedReader) Read(p []byte) (n int, err error) {
return r.buffer.Read(p)
}
func (r *BufferedReader) WriteTo(w io.Writer) (n int64, err error) {
return CopyExtendedBuffer(NewExtendedWriter(w), NewExtendedReader(r.upstream), r.buffer)
}
func (w *BufferedReader) Upstream() any {
return w.upstream
}

31
common/bufio/conn.go
View file

@ -4,6 +4,7 @@ import (
"context"
"io"
"net"
"os"
"time"
"github.com/sagernet/sing/common"
@ -14,6 +15,27 @@ import (
"github.com/sagernet/sing/common/task"
)
type readOnlyReader struct {
io.Reader
}
func (r *readOnlyReader) WriteTo(w io.Writer) (n int64, err error) {
return Copy(w, r.Reader)
}
func needReadFromWrapper(dst io.ReaderFrom, src io.Reader) bool {
_, isTCPConn := dst.(*net.TCPConn)
if !isTCPConn {
return false
}
switch src.(type) {
case *net.TCPConn, *net.UnixConn, *os.File:
return false
default:
return true
}
}
func Copy(dst io.Writer, src io.Reader) (n int64, err error) {
src = N.UnwrapReader(src)
dst = N.UnwrapWriter(dst)
@ -21,6 +43,9 @@ func Copy(dst io.Writer, src io.Reader) (n int64, err error) {
return wt.WriteTo(dst)
}
if rt, ok := dst.(io.ReaderFrom); ok {
if needReadFromWrapper(rt, src) {
src = &readOnlyReader{src}
}
return rt.ReadFrom(src)
}
return CopyExtended(NewExtendedWriter(dst), NewExtendedReader(src))
@ -37,7 +62,11 @@ func CopyExtended(dst N.ExtendedWriter, src N.ExtendedReader) (n int64, err erro
_buffer := buf.StackNew()
defer common.KeepAlive(_buffer)
buffer := common.Dup(_buffer)
return CopyExtendedBuffer(dst, src, common.Dup(_buffer))
}
func CopyExtendedBuffer(dst N.ExtendedWriter, src N.ExtendedReader, buffer *buf.Buffer) (n int64, err error) {
buffer.IncRef()
defer buffer.DecRef()
for {