diff --git a/common/buf/io.go b/common/buf/io.go index aedfcbf27..bed9d84f0 100644 --- a/common/buf/io.go +++ b/common/buf/io.go @@ -92,7 +92,7 @@ func NewWriter(writer io.Writer) Writer { // NewSequentialWriter returns a Writer that write Buffers in a MultiBuffer sequentially. func NewSequentialWriter(writer io.Writer) Writer { - return &seqWriter{ - writer: writer, + return &SequentialWriter{ + Writer: writer, } } diff --git a/common/buf/reader.go b/common/buf/reader.go index cd9d93a8f..59d48709c 100644 --- a/common/buf/reader.go +++ b/common/buf/reader.go @@ -7,6 +7,22 @@ import ( "v2ray.com/core/common/errors" ) +func readOne(r io.Reader) (*Buffer, error) { + b := New() + for i := 0; i < 64; i++ { + err := b.Reset(ReadFrom(r)) + if !b.IsEmpty() { + return b, nil + } + if err != nil { + b.Release() + return nil, err + } + } + + return nil, newError("Reader returns too many empty payloads.") +} + // BytesToBufferReader is a Reader that adjusts its reading speed automatically. type BytesToBufferReader struct { io.Reader @@ -21,22 +37,14 @@ func NewBytesToBufferReader(reader io.Reader) Reader { } func (r *BytesToBufferReader) readSmall() (MultiBuffer, error) { - b := New() - for i := 0; i < 64; i++ { - err := b.Reset(ReadFrom(r.Reader)) - if b.IsFull() && largeSize > Size { - r.buffer = newBytes(Size + 1) - } - if !b.IsEmpty() { - return NewMultiBufferValue(b), nil - } - if err != nil { - b.Release() - return nil, err - } + b, err := readOne(r.Reader) + if b.IsFull() && largeSize > Size { + r.buffer = newBytes(Size + 1) } - - return nil, newError("Reader returns too many empty payloads.") + if err != nil { + return nil, err + } + return NewMultiBufferValue(b), nil } func (r *BytesToBufferReader) freeBuffer() { @@ -195,3 +203,15 @@ func (r *BufferedReader) Close() error { } return common.Close(r.Reader) } + +type SingleReader struct { + io.Reader +} + +func (r *SingleReader) ReadMultiBuffer() (MultiBuffer, error) { + b, err := readOne(r.Reader) + if err != nil { + return nil, err + } + return NewMultiBufferValue(b), nil +} diff --git a/common/buf/writer.go b/common/buf/writer.go index b6b9e0e94..04ca3da42 100644 --- a/common/buf/writer.go +++ b/common/buf/writer.go @@ -171,15 +171,15 @@ func (w *BufferedWriter) Close() error { return common.Close(w.writer) } -type seqWriter struct { - writer io.Writer +type SequentialWriter struct { + io.Writer } -func (w *seqWriter) WriteMultiBuffer(mb MultiBuffer) error { +func (w *SequentialWriter) WriteMultiBuffer(mb MultiBuffer) error { defer mb.Release() for _, b := range mb { - if err := WriteAllBytes(w.writer, b.Bytes()); err != nil { + if err := WriteAllBytes(w.Writer, b.Bytes()); err != nil { return err } } diff --git a/common/crypto/auth.go b/common/crypto/auth.go index 4cbfeee8e..cc623e52c 100644 --- a/common/crypto/auth.go +++ b/common/crypto/auth.go @@ -172,10 +172,11 @@ func (r *AuthenticationReader) ReadMultiBuffer() (buf.MultiBuffer, error) { return nil, err } - mb := buf.NewMultiBufferCap(32) + const readSize = 16 + mb := buf.NewMultiBufferCap(readSize) mb.Append(b) - for { + for i := 1; i < readSize; i++ { b, err := r.readInternal(true) if err == errSoft || err == io.EOF { break @@ -288,6 +289,7 @@ func (w *AuthenticationWriter) writePacket(mb buf.MultiBuffer) error { return w.writer.WriteMultiBuffer(mb2Write) } +// WriteMultiBuffer implements buf.Writer. func (w *AuthenticationWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { if mb.IsEmpty() { b := buf.New() diff --git a/proxy/shadowsocks/config.go b/proxy/shadowsocks/config.go index 9735ee874..df720f8f4 100644 --- a/proxy/shadowsocks/config.go +++ b/proxy/shadowsocks/config.go @@ -124,12 +124,14 @@ func (v *AesCfb) IVSize() int32 { func (v *AesCfb) NewEncryptionWriter(key []byte, iv []byte, writer io.Writer) (buf.Writer, error) { stream := crypto.NewAesEncryptionStream(key, iv) - return buf.NewWriter(crypto.NewCryptionWriter(stream, writer)), nil + return &buf.SequentialWriter{Writer: crypto.NewCryptionWriter(stream, writer)}, nil } func (v *AesCfb) NewDecryptionReader(key []byte, iv []byte, reader io.Reader) (buf.Reader, error) { stream := crypto.NewAesDecryptionStream(key, iv) - return buf.NewReader(crypto.NewCryptionReader(stream, reader)), nil + return &buf.SingleReader{ + Reader: crypto.NewCryptionReader(stream, reader), + }, nil } func (v *AesCfb) EncodePacket(key []byte, b *buf.Buffer) error { @@ -243,12 +245,12 @@ func (v *ChaCha20) IVSize() int32 { func (v *ChaCha20) NewEncryptionWriter(key []byte, iv []byte, writer io.Writer) (buf.Writer, error) { stream := crypto.NewChaCha20Stream(key, iv) - return buf.NewWriter(crypto.NewCryptionWriter(stream, writer)), nil + return &buf.SequentialWriter{Writer: crypto.NewCryptionWriter(stream, writer)}, nil } func (v *ChaCha20) NewDecryptionReader(key []byte, iv []byte, reader io.Reader) (buf.Reader, error) { stream := crypto.NewChaCha20Stream(key, iv) - return buf.NewReader(crypto.NewCryptionReader(stream, reader)), nil + return &buf.SingleReader{Reader: crypto.NewCryptionReader(stream, reader)}, nil } func (v *ChaCha20) EncodePacket(key []byte, b *buf.Buffer) error {