Explorar el Código

optimize encryption read/write operations

Darien Raymond hace 7 años
padre
commit
6a06908456
Se han modificado 5 ficheros con 51 adiciones y 27 borrados
  1. 2 2
      common/buf/io.go
  2. 35 15
      common/buf/reader.go
  3. 4 4
      common/buf/writer.go
  4. 4 2
      common/crypto/auth.go
  5. 6 4
      proxy/shadowsocks/config.go

+ 2 - 2
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,
 	}
 }

+ 35 - 15
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
+}

+ 4 - 4
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
 		}
 	}

+ 4 - 2
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()

+ 6 - 4
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 {