Fix: CounterConnection with ReadV/WriteV (#720)

Co-authored-by: JimhHan <50871214+JimhHan@users.noreply.github.com>
This commit is contained in:
Arthur Morgan 2021-09-20 20:11:21 +08:00 committed by GitHub
parent f2cb13a8ec
commit 24b637cd5e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
53 changed files with 247 additions and 128 deletions

View file

@ -5,6 +5,8 @@ import (
"net"
"sync"
"github.com/xtls/xray-core/features/stats"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/errors"
)
@ -13,7 +15,8 @@ import (
type BufferToBytesWriter struct {
io.Writer
cache [][]byte
counter stats.Counter
cache [][]byte
}
// WriteMultiBuffer implements Writer. This method takes ownership of the given buffer.
@ -26,7 +29,7 @@ func (w *BufferToBytesWriter) WriteMultiBuffer(mb MultiBuffer) error {
}
if len(mb) == 1 {
return WriteAllBytes(w.Writer, mb[0].Bytes())
return WriteAllBytes(w.Writer, mb[0].Bytes(), w.counter)
}
if cap(w.cache) < len(mb) {
@ -45,9 +48,15 @@ func (w *BufferToBytesWriter) WriteMultiBuffer(mb MultiBuffer) error {
}()
nb := net.Buffers(bs)
wc := int64(0)
defer func() {
if w.counter != nil {
w.counter.Add(wc)
}
}()
for size > 0 {
n, err := nb.WriteTo(w.Writer)
wc += n
if err != nil {
return err
}
@ -173,7 +182,7 @@ func (w *BufferedWriter) flushInternal() error {
w.buffer = nil
if writer, ok := w.writer.(io.Writer); ok {
err := WriteAllBytes(writer, b.Bytes())
err := WriteAllBytes(writer, b.Bytes(), nil)
b.Release()
return err
}