Added TestsEnabled Settings to enable VMessAEAD test

chyspace · Jun 8, 2020 · d06a4d1 · d06a4d1
1 parent 00103d1
commit d06a4d1
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -0,0 +1,3 @@
+package vmess
+
+const TestsEnabled = "VMessCtxInterface_TestsEnabled"
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		package vmess

		const TestsEnabled = "VMessCtxInterface_TestsEnabled"