mirror of
https://github.com/XTLS/Xray-core.git
synced 2025-04-30 09:18:34 +00:00
Fix: CounterConnection with ReadV/WriteV (#720)
Co-authored-by: JimhHan <50871214+JimhHan@users.noreply.github.com>
This commit is contained in:
parent
f2cb13a8ec
commit
24b637cd5e
53 changed files with 247 additions and 128 deletions
|
@ -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
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue