mirror of https://github.com/XTLS/Xray-core
156 lines
2.6 KiB
Go
156 lines
2.6 KiB
Go
//go:build !wasm
|
|
// +build !wasm
|
|
|
|
package buf
|
|
|
|
import (
|
|
"io"
|
|
"syscall"
|
|
|
|
"github.com/xtls/xray-core/common/platform"
|
|
"github.com/xtls/xray-core/features/stats"
|
|
)
|
|
|
|
type allocStrategy struct {
|
|
current uint32
|
|
}
|
|
|
|
func (s *allocStrategy) Current() uint32 {
|
|
return s.current
|
|
}
|
|
|
|
func (s *allocStrategy) Adjust(n uint32) {
|
|
if n >= s.current {
|
|
s.current *= 2
|
|
} else {
|
|
s.current = n
|
|
}
|
|
|
|
if s.current > 8 {
|
|
s.current = 8
|
|
}
|
|
|
|
if s.current == 0 {
|
|
s.current = 1
|
|
}
|
|
}
|
|
|
|
func (s *allocStrategy) Alloc() []*Buffer {
|
|
bs := make([]*Buffer, s.current)
|
|
for i := range bs {
|
|
bs[i] = New()
|
|
}
|
|
return bs
|
|
}
|
|
|
|
type multiReader interface {
|
|
Init([]*Buffer)
|
|
Read(fd uintptr) int32
|
|
Clear()
|
|
}
|
|
|
|
// ReadVReader is a Reader that uses readv(2) syscall to read data.
|
|
type ReadVReader struct {
|
|
io.Reader
|
|
rawConn syscall.RawConn
|
|
mr multiReader
|
|
alloc allocStrategy
|
|
counter stats.Counter
|
|
}
|
|
|
|
// NewReadVReader creates a new ReadVReader.
|
|
func NewReadVReader(reader io.Reader, rawConn syscall.RawConn, counter stats.Counter) *ReadVReader {
|
|
return &ReadVReader{
|
|
Reader: reader,
|
|
rawConn: rawConn,
|
|
alloc: allocStrategy{
|
|
current: 1,
|
|
},
|
|
mr: newMultiReader(),
|
|
counter: counter,
|
|
}
|
|
}
|
|
|
|
func (r *ReadVReader) readMulti() (MultiBuffer, error) {
|
|
bs := r.alloc.Alloc()
|
|
|
|
r.mr.Init(bs)
|
|
var nBytes int32
|
|
err := r.rawConn.Read(func(fd uintptr) bool {
|
|
n := r.mr.Read(fd)
|
|
if n < 0 {
|
|
return false
|
|
}
|
|
|
|
nBytes = n
|
|
return true
|
|
})
|
|
r.mr.Clear()
|
|
|
|
if err != nil {
|
|
ReleaseMulti(MultiBuffer(bs))
|
|
return nil, err
|
|
}
|
|
|
|
if nBytes == 0 {
|
|
ReleaseMulti(MultiBuffer(bs))
|
|
return nil, io.EOF
|
|
}
|
|
|
|
nBuf := 0
|
|
for nBuf < len(bs) {
|
|
if nBytes <= 0 {
|
|
break
|
|
}
|
|
end := nBytes
|
|
if end > Size {
|
|
end = Size
|
|
}
|
|
bs[nBuf].end = end
|
|
nBytes -= end
|
|
nBuf++
|
|
}
|
|
|
|
for i := nBuf; i < len(bs); i++ {
|
|
bs[i].Release()
|
|
bs[i] = nil
|
|
}
|
|
|
|
return MultiBuffer(bs[:nBuf]), nil
|
|
}
|
|
|
|
// ReadMultiBuffer implements Reader.
|
|
func (r *ReadVReader) ReadMultiBuffer() (MultiBuffer, error) {
|
|
if r.alloc.Current() == 1 {
|
|
b, err := ReadBuffer(r.Reader)
|
|
if b.IsFull() {
|
|
r.alloc.Adjust(1)
|
|
}
|
|
if r.counter != nil && b != nil {
|
|
r.counter.Add(int64(b.Len()))
|
|
}
|
|
return MultiBuffer{b}, err
|
|
}
|
|
|
|
mb, err := r.readMulti()
|
|
if r.counter != nil && mb != nil {
|
|
r.counter.Add(int64(mb.Len()))
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
r.alloc.Adjust(uint32(len(mb)))
|
|
return mb, nil
|
|
}
|
|
|
|
var useReadv bool
|
|
|
|
func init() {
|
|
const defaultFlagValue = "NOT_DEFINED_AT_ALL"
|
|
value := platform.NewEnvFlag(platform.UseReadV).GetValue(func() string { return defaultFlagValue })
|
|
switch value {
|
|
case defaultFlagValue, "auto", "enable":
|
|
useReadv = true
|
|
}
|
|
}
|