From d06a4d1f505fbe7402a0c127de819cf62f49bd1c Mon Sep 17 00:00:00 2001
From: Shelikhoo <xiaokangwang@outlook.com>
Date: Mon, 8 Jun 2020 16:50:50 +0800
Subject: [PATCH] Added TestsEnabled Settings to enable VMessAEAD test

---
 infra/conf/vmess.go                   |  8 +++++---
 proxy/vmess/account.go                |  9 ++++++---
 proxy/vmess/encoding/client.go        | 13 +++++++++++--
 proxy/vmess/encoding/encoding_test.go |  7 ++++---
 proxy/vmess/outbound/outbound.go      |  7 +++++--
 proxy/vmess/vmessCtxInterface.go      |  3 +++
 6 files changed, 34 insertions(+), 13 deletions(-)
 create mode 100644 proxy/vmess/vmessCtxInterface.go

diff --git a/infra/conf/vmess.go b/infra/conf/vmess.go
index d243305aa..b9531fdec 100644
--- a/infra/conf/vmess.go
+++ b/infra/conf/vmess.go
@@ -14,9 +14,10 @@ import (
 )
 
 type VMessAccount struct {
-	ID       string `json:"id"`
-	AlterIds uint16 `json:"alterId"`
-	Security string `json:"security"`
+	ID           string `json:"id"`
+	AlterIds     uint16 `json:"alterId"`
+	Security     string `json:"security"`
+	TestsEnabled string `json:"testsEnabled"`
 }
 
 // Build implements Buildable
@@ -40,6 +41,7 @@ func (a *VMessAccount) Build() *vmess.Account {
 		SecuritySettings: &protocol.SecurityConfig{
 			Type: st,
 		},
+		TestsEnabled: a.TestsEnabled,
 	}
 }
 
diff --git a/proxy/vmess/account.go b/proxy/vmess/account.go
index cdced1539..8e621f8d8 100644
--- a/proxy/vmess/account.go
+++ b/proxy/vmess/account.go
@@ -16,6 +16,8 @@ type MemoryAccount struct {
 	AlterIDs []*protocol.ID
 	// Security type of the account. Used for client connections.
 	Security protocol.SecurityType
+
+	TestsEnabled string
 }
 
 // AnyValidID returns an ID that is either the main ID or one of the alternative IDs if any.
@@ -44,8 +46,9 @@ func (a *Account) AsAccount() (protocol.Account, error) {
 	}
 	protoID := protocol.NewID(id)
 	return &MemoryAccount{
-		ID:       protoID,
-		AlterIDs: protocol.NewAlterIDs(protoID, uint16(a.AlterId)),
-		Security: a.SecuritySettings.GetSecurityType(),
+		ID:           protoID,
+		AlterIDs:     protocol.NewAlterIDs(protoID, uint16(a.AlterId)),
+		Security:     a.SecuritySettings.GetSecurityType(),
+		TestsEnabled: a.TestsEnabled,
 	}, nil
 }
diff --git a/proxy/vmess/encoding/client.go b/proxy/vmess/encoding/client.go
index d34ef484c..9ecd74824 100644
--- a/proxy/vmess/encoding/client.go
+++ b/proxy/vmess/encoding/client.go
@@ -2,6 +2,7 @@ package encoding
 
 import (
 	"bytes"
+	"context"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/md5"
@@ -13,6 +14,7 @@ import (
 	"hash/fnv"
 	"io"
 	"os"
+	"strings"
 	vmessaead "v2ray.com/core/proxy/vmess/aead"
 
 	"golang.org/x/crypto/chacha20poly1305"
@@ -49,13 +51,20 @@ type ClientSession struct {
 }
 
 // NewClientSession creates a new ClientSession.
-func NewClientSession(idHash protocol.IDHash) *ClientSession {
+func NewClientSession(idHash protocol.IDHash, ctx context.Context) *ClientSession {
 	randomBytes := make([]byte, 33) // 16 + 16 + 1
 	common.Must2(rand.Read(randomBytes))
 
 	session := &ClientSession{}
 
-	session.isAEADRequest = true
+	session.isAEADRequest = false
+
+	if ctxValueTestsEnabled := ctx.Value(vmess.TestsEnabled); ctxValueTestsEnabled != nil {
+		testsEnabled := ctxValueTestsEnabled.(string)
+		if strings.Contains(testsEnabled, "VMessAEAD") {
+			session.isAEADRequest = true
+		}
+	}
 
 	if vmessexp, vmessexp_found := os.LookupEnv("VMESSAEADEXPERIMENT"); vmessexp_found {
 		if vmessexp == "y" {
diff --git a/proxy/vmess/encoding/encoding_test.go b/proxy/vmess/encoding/encoding_test.go
index 00fef1ecf..bc7eecd33 100644
--- a/proxy/vmess/encoding/encoding_test.go
+++ b/proxy/vmess/encoding/encoding_test.go
@@ -1,6 +1,7 @@
 package encoding_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -42,7 +43,7 @@ func TestRequestSerialization(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(protocol.DefaultIDHash)
+	client := NewClientSession(protocol.DefaultIDHash, context.TODO())
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()
@@ -92,7 +93,7 @@ func TestInvalidRequest(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(protocol.DefaultIDHash)
+	client := NewClientSession(protocol.DefaultIDHash, context.TODO())
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()
@@ -133,7 +134,7 @@ func TestMuxRequest(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(protocol.DefaultIDHash)
+	client := NewClientSession(protocol.DefaultIDHash, context.TODO())
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()
diff --git a/proxy/vmess/outbound/outbound.go b/proxy/vmess/outbound/outbound.go
index 91c80b4f8..5ab123bb0 100644
--- a/proxy/vmess/outbound/outbound.go
+++ b/proxy/vmess/outbound/outbound.go
@@ -89,9 +89,10 @@ func (v *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 		command = protocol.RequestCommandMux
 	}
 
+	user := rec.PickUser()
 	request := &protocol.RequestHeader{
 		Version: encoding.Version,
-		User:    rec.PickUser(),
+		User:    user,
 		Command: command,
 		Address: target.Address,
 		Port:    target.Port,
@@ -112,7 +113,9 @@ func (v *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 	input := link.Reader
 	output := link.Writer
 
-	session := encoding.NewClientSession(protocol.DefaultIDHash)
+	ctx = context.WithValue(ctx, vmess.TestsEnabled, user.Account.(*vmess.MemoryAccount).TestsEnabled)
+
+	session := encoding.NewClientSession(protocol.DefaultIDHash, ctx)
 	sessionPolicy := v.policyManager.ForLevel(request.User.Level)
 
 	ctx, cancel := context.WithCancel(ctx)
diff --git a/proxy/vmess/vmessCtxInterface.go b/proxy/vmess/vmessCtxInterface.go
new file mode 100644
index 000000000..edaad0606
--- /dev/null
+++ b/proxy/vmess/vmessCtxInterface.go
@@ -0,0 +1,3 @@
+package vmess
+
+const TestsEnabled = "VMessCtxInterface_TestsEnabled"