From 0bc846f016c5ea4c7974f39ef10766943925447e Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Thu, 5 May 2016 00:24:18 +0200
Subject: [PATCH] Fix ending in vmess reader.

---
 proxy/dokodemo/dokodemo.go              |  3 +++
 proxy/vmess/io/reader.go                | 12 ++++++------
 testing/scenarios/dynamic_vmess_test.go |  8 +++++---
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index 57b33f3df..c1b2a0e5d 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -129,6 +129,8 @@ func (this *DokodemoDoor) HandleTCPConnection(conn *hub.Connection) {
 	outputFinish.Lock()
 
 	reader := v2net.NewTimeOutReader(this.config.Timeout, conn)
+	defer reader.Release()
+
 	go func() {
 		v2reader := v2io.NewAdaptiveReader(reader)
 		defer v2reader.Release()
@@ -147,4 +149,5 @@ func (this *DokodemoDoor) HandleTCPConnection(conn *hub.Connection) {
 	}()
 
 	outputFinish.Lock()
+	inputFinish.Lock()
 }
diff --git a/proxy/vmess/io/reader.go b/proxy/vmess/io/reader.go
index a9debadf9..27e622af4 100644
--- a/proxy/vmess/io/reader.go
+++ b/proxy/vmess/io/reader.go
@@ -62,12 +62,6 @@ func (this *AuthChunkReader) Read() (*alloc.Buffer, error) {
 		buffer = AllocBuffer(this.chunkLength).Clear()
 	}
 
-	_, err := buffer.FillFrom(this.reader)
-	if err != nil {
-		buffer.Release()
-		return nil, err
-	}
-
 	if this.chunkLength == -1 {
 		for buffer.Len() < 6 {
 			_, err := buffer.FillFrom(this.reader)
@@ -80,6 +74,12 @@ func (this *AuthChunkReader) Read() (*alloc.Buffer, error) {
 		this.chunkLength = int(length) - 4
 		this.validator = NewValidator(serial.BytesLiteral(buffer.Value[2:6]).Uint32Value())
 		buffer.SliceFrom(6)
+	} else if buffer.Len() < this.chunkLength {
+		_, err := buffer.FillFrom(this.reader)
+		if err != nil {
+			buffer.Release()
+			return nil, err
+		}
 	}
 
 	if this.chunkLength == 0 {
diff --git a/testing/scenarios/dynamic_vmess_test.go b/testing/scenarios/dynamic_vmess_test.go
index 6dbbb64fc..c61c8cf84 100644
--- a/testing/scenarios/dynamic_vmess_test.go
+++ b/testing/scenarios/dynamic_vmess_test.go
@@ -1,6 +1,8 @@
 package scenarios
 
 import (
+	"bytes"
+	"io"
 	"net"
 	"testing"
 
@@ -42,10 +44,10 @@ func TestDynamicVMess(t *testing.T) {
 
 		conn.CloseWrite()
 
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
+		response := bytes.NewBuffer(nil)
+		_, err = io.Copy(response, conn)
 		assert.Error(err).IsNil()
-		assert.StringLiteral("Processed: " + payload).Equals(string(response[:nBytes]))
+		assert.StringLiteral("Processed: " + payload).Equals(string(response.Bytes()))
 
 		conn.Close()
 	}