Unified drain support for vmess and shadowsocks (#791)

* Added test for no terminate signal

* unified drain support for vmess and shadowsockets

* drain: add generated file

Co-authored-by: Shelikhoo <xiaokangwang@outlook.com>

common/drain/drain.go

@@ -0,0 +1,10 @@
+package drain
+
+import "io"
+
+//go:generate go run github.com/xtls/xray-core/common/errors/errorgen
+
+type Drainer interface {
+	AcknowledgeReceive(size int)
+	Drain(reader io.Reader) error
+}

common/drain/drainer.go

@@ -0,0 +1,62 @@
+package drain
+
+import (
+	"io"
+	"io/ioutil"
+
+	"github.com/xtls/xray-core/common/dice"
+)
+
+type BehaviorSeedLimitedDrainer struct {
+	DrainSize int
+}
+
+func NewBehaviorSeedLimitedDrainer(behaviorSeed int64, drainFoundation, maxBaseDrainSize, maxRandDrain int) (Drainer, error) {
+	behaviorRand := dice.NewDeterministicDice(behaviorSeed)
+	BaseDrainSize := behaviorRand.Roll(maxBaseDrainSize)
+	RandDrainMax := behaviorRand.Roll(maxRandDrain) + 1
+	RandDrainRolled := dice.Roll(RandDrainMax)
+	DrainSize := drainFoundation + BaseDrainSize + RandDrainRolled
+	return &BehaviorSeedLimitedDrainer{DrainSize: DrainSize}, nil
+}
+
+func (d *BehaviorSeedLimitedDrainer) AcknowledgeReceive(size int) {
+	d.DrainSize -= size
+}
+
+func (d *BehaviorSeedLimitedDrainer) Drain(reader io.Reader) error {
+	if d.DrainSize > 0 {
+		err := drainReadN(reader, d.DrainSize)
+		if err == nil {
+			return newError("drained connection")
+		}
+		return newError("unable to drain connection").Base(err)
+	}
+	return nil
+}
+
+func drainReadN(reader io.Reader, n int) error {
+	_, err := io.CopyN(ioutil.Discard, reader, int64(n))
+	return err
+}
+
+func WithError(drainer Drainer, reader io.Reader, err error) error {
+	drainErr := drainer.Drain(reader)
+	if drainErr == nil {
+		return err
+	}
+	return newError(drainErr).Base(err)
+}
+
+type NopDrainer struct{}
+
+func (n NopDrainer) AcknowledgeReceive(size int) {
+}
+
+func (n NopDrainer) Drain(reader io.Reader) error {
+	return nil
+}
+
+func NewNopDrainer() Drainer {
+	return &NopDrainer{}
+}

common/drain/errors.generated.go

@@ -0,0 +1,9 @@
+package drain
+
+import "github.com/xtls/xray-core/common/errors"
+
+type errPathObjHolder struct{}
+
+func newError(values ...interface{}) *errors.Error {
+	return errors.New(values...).WithPathObj(errPathObjHolder{})
+}

proxy/shadowsocks/protocol.go

@@ -1,13 +1,16 @@
 package shadowsocks
 
 import (
+	"crypto/hmac"
 	"crypto/rand"
+	"crypto/sha256"
+	"hash/crc32"
 	"io"
 
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/crypto"
-	"github.com/xtls/xray-core/common/dice"
+	"github.com/xtls/xray-core/common/drain"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/protocol"
 )
@@ -51,21 +54,19 @@ func (r *FullReader) Read(p []byte) (n int, err error) {
 // ReadTCPSession reads a Shadowsocks TCP session from the given reader, returns its header and remaining parts.
 func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHeader, buf.Reader, error) {
 	behaviorSeed := validator.GetBehaviorSeed()
-	behaviorRand := dice.NewDeterministicDice(int64(behaviorSeed))
-	BaseDrainSize := behaviorRand.Roll(3266)
-	RandDrainMax := behaviorRand.Roll(64) + 1
-	RandDrainRolled := dice.Roll(RandDrainMax)
-	DrainSize := BaseDrainSize + 16 + 38 + RandDrainRolled
-	readSizeRemain := DrainSize
+	drainer, errDrain := drain.NewBehaviorSeedLimitedDrainer(int64(behaviorSeed), 16+38, 3266, 64)
+
+	if errDrain != nil {
+		return nil, nil, newError("failed to initialize drainer").Base(errDrain)
+	}
 
 	var r buf.Reader
 	buffer := buf.New()
 	defer buffer.Release()
 
 	if _, err := buffer.ReadFullFrom(reader, 50); err != nil {
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("failed to read 50 bytes").Base(err)
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		return nil, nil, drain.WithError(drainer, reader, newError("failed to read 50 bytes").Base(err))
 	}
 
 	bs := buffer.Bytes()
@@ -73,16 +74,14 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 
 	switch err {
 	case ErrNotFound:
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("failed to match an user").Base(err)
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		return nil, nil, drain.WithError(drainer, reader, newError("failed to match an user").Base(err))
 	case ErrIVNotUnique:
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("failed iv check").Base(err)
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		return nil, nil, drain.WithError(drainer, reader, newError("failed iv check").Base(err))
 	default:
 		reader = &FullReader{reader, bs[ivLen:]}
-		readSizeRemain -= int(ivLen)
+		drainer.AcknowledgeReceive(int(ivLen))
 
 		if aead != nil {
 			auth := &crypto.AEADAuthenticator{
@@ -97,8 +96,7 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 			iv := append([]byte(nil), buffer.BytesTo(ivLen)...)
 			r, err = account.Cipher.NewDecryptionReader(account.Key, iv, reader)
 			if err != nil {
-				DrainConnN(reader, readSizeRemain)
-				return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
+				return nil, nil, drain.WithError(drainer, reader, newError("failed to initialize decoding stream").Base(err).AtError())
 			}
 		}
 	}
@@ -115,28 +113,21 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 
 	addr, port, err := addrParser.ReadAddressPort(buffer, br)
 	if err != nil {
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("failed to read address").Base(err)
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		return nil, nil, drain.WithError(drainer, reader, newError("failed to read address").Base(err))
 	}
 
 	request.Address = addr
 	request.Port = port
 
 	if request.Address == nil {
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("invalid remote address.")
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		return nil, nil, drain.WithError(drainer, reader, newError("invalid remote address."))
 	}
 
 	return request, br, nil
 }
 
-func DrainConnN(reader io.Reader, n int) error {
-	_, err := io.CopyN(io.Discard, reader, int64(n))
-	return err
-}
-
 // WriteTCPRequest writes Shadowsocks request into the given writer, and returns a writer for body.
 func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Writer, error) {
 	user := request.User
@@ -175,16 +166,29 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Wri
 func ReadTCPResponse(user *protocol.MemoryUser, reader io.Reader) (buf.Reader, error) {
 	account := user.Account.(*MemoryAccount)
 
+	hashkdf := hmac.New(sha256.New, []byte("SSBSKDF"))
+	hashkdf.Write(account.Key)
+
+	behaviorSeed := crc32.ChecksumIEEE(hashkdf.Sum(nil))
+
+	drainer, err := drain.NewBehaviorSeedLimitedDrainer(int64(behaviorSeed), 16+38, 3266, 64)
+
+	if err != nil {
+		return nil, newError("failed to initialize drainer").Base(err)
+	}
+
 	var iv []byte
 	if account.Cipher.IVSize() > 0 {
 		iv = make([]byte, account.Cipher.IVSize())
-		if _, err := io.ReadFull(reader, iv); err != nil {
+		if n, err := io.ReadFull(reader, iv); err != nil {
 			return nil, newError("failed to read IV").Base(err)
+		} else { // nolint: golint
+			drainer.AcknowledgeReceive(n)
 		}
 	}
 
 	if ivError := account.CheckIV(iv); ivError != nil {
-		return nil, newError("failed iv check").Base(ivError)
+		return nil, drain.WithError(drainer, reader, newError("failed iv check").Base(ivError))
 	}
 
 	return account.Cipher.NewDecryptionReader(account.Key, iv, reader)

proxy/vmess/encoding/client.go

@@ -13,6 +13,8 @@ import (
 	"hash/fnv"
 	"io"
 
+	"github.com/xtls/xray-core/common/drain"
+
 	"golang.org/x/crypto/chacha20poly1305"
 
 	"github.com/xtls/xray-core/common"
@@ -44,10 +46,12 @@ type ClientSession struct {
 	responseBodyIV  [16]byte
 	responseReader  io.Reader
 	responseHeader  byte
+
+	readDrainer drain.Drainer
 }
 
 // NewClientSession creates a new ClientSession.
-func NewClientSession(ctx context.Context, isAEAD bool, idHash protocol.IDHash) *ClientSession {
+func NewClientSession(ctx context.Context, isAEAD bool, idHash protocol.IDHash, behaviorSeed int64) *ClientSession {
 	session := &ClientSession{
 		isAEAD: isAEAD,
 		idHash: idHash,
@@ -68,6 +72,14 @@ func NewClientSession(ctx context.Context, isAEAD bool, idHash protocol.IDHash)
 		BodyIV := sha256.Sum256(session.requestBodyIV[:])
 		copy(session.responseBodyIV[:], BodyIV[:16])
 	}
+	{
+		var err error
+		session.readDrainer, err = drain.NewBehaviorSeedLimitedDrainer(behaviorSeed, 18, 3266, 64)
+		if err != nil {
+			newError("unable to initialize drainer").Base(err).WriteToLog()
+			session.readDrainer = drain.NewNopDrainer()
+		}
+	}
 
 	return session
 }
@@ -225,12 +237,15 @@ func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.Respon
 		var decryptedResponseHeaderLength int
 		var decryptedResponseHeaderLengthBinaryDeserializeBuffer uint16
 
-		if _, err := io.ReadFull(reader, aeadEncryptedResponseHeaderLength[:]); err != nil {
-			return nil, newError("Unable to Read Header Len").Base(err)
+		if n, err := io.ReadFull(reader, aeadEncryptedResponseHeaderLength[:]); err != nil {
+			c.readDrainer.AcknowledgeReceive(n)
+			return nil, drain.WithError(c.readDrainer, reader, newError("Unable to Read Header Len").Base(err))
+		} else { // nolint: golint
+			c.readDrainer.AcknowledgeReceive(n)
 		}
 		if decryptedResponseHeaderLengthBinaryBuffer, err := aeadResponseHeaderLengthEncryptionAEAD.Open(nil, aeadResponseHeaderLengthEncryptionIV, aeadEncryptedResponseHeaderLength[:], nil); err != nil {
-			return nil, newError("Failed To Decrypt Length").Base(err)
-		} else {
+			return nil, drain.WithError(c.readDrainer, reader, newError("Failed To Decrypt Length").Base(err))
+		} else { // nolint: golint
 			common.Must(binary.Read(bytes.NewReader(decryptedResponseHeaderLengthBinaryBuffer), binary.BigEndian, &decryptedResponseHeaderLengthBinaryDeserializeBuffer))
 			decryptedResponseHeaderLength = int(decryptedResponseHeaderLengthBinaryDeserializeBuffer)
 		}
@@ -243,13 +258,16 @@ func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.Respon
 
 		encryptedResponseHeaderBuffer := make([]byte, decryptedResponseHeaderLength+16)
 
-		if _, err := io.ReadFull(reader, encryptedResponseHeaderBuffer); err != nil {
-			return nil, newError("Unable to Read Header Data").Base(err)
+		if n, err := io.ReadFull(reader, encryptedResponseHeaderBuffer); err != nil {
+			c.readDrainer.AcknowledgeReceive(n)
+			return nil, drain.WithError(c.readDrainer, reader, newError("Unable to Read Header Data").Base(err))
+		} else { // nolint: golint
+			c.readDrainer.AcknowledgeReceive(n)
 		}
 
 		if decryptedResponseHeaderBuffer, err := aeadResponseHeaderPayloadEncryptionAEAD.Open(nil, aeadResponseHeaderPayloadEncryptionIV, encryptedResponseHeaderBuffer, nil); err != nil {
-			return nil, newError("Failed To Decrypt Payload").Base(err)
-		} else {
+			return nil, drain.WithError(c.readDrainer, reader, newError("Failed To Decrypt Payload").Base(err))
+		} else { // nolint: golint
 			c.responseReader = bytes.NewReader(decryptedResponseHeaderBuffer)
 		}
 	}

proxy/vmess/encoding/encoding_test.go

@@ -43,7 +43,7 @@ func TestRequestSerialization(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash)
+	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash, 0)
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()
@@ -93,7 +93,7 @@ func TestInvalidRequest(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash)
+	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash, 0)
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()
@@ -134,7 +134,7 @@ func TestMuxRequest(t *testing.T) {
 	}
 
 	buffer := buf.New()
-	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash)
+	client := NewClientSession(context.TODO(), true, protocol.DefaultIDHash, 0)
 	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))
 
 	buffer2 := buf.New()

proxy/vmess/encoding/server.go

@@ -12,13 +12,14 @@ import (
 	"sync"
 	"time"
 
+	"github.com/xtls/xray-core/common/drain"
+
 	"golang.org/x/crypto/chacha20poly1305"
 
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/bitmask"
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/crypto"
-	"github.com/xtls/xray-core/common/dice"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/protocol"
 	"github.com/xtls/xray-core/common/task"
@@ -138,22 +139,18 @@ func parseSecurityType(b byte) protocol.SecurityType {
 // DecodeRequestHeader decodes and returns (if successful) a RequestHeader from an input stream.
 func (s *ServerSession) DecodeRequestHeader(reader io.Reader, isDrain bool) (*protocol.RequestHeader, error) {
 	buffer := buf.New()
-	behaviorRand := dice.NewDeterministicDice(int64(s.userValidator.GetBehaviorSeed()))
-	BaseDrainSize := behaviorRand.Roll(3266)
-	RandDrainMax := behaviorRand.Roll(64) + 1
-	RandDrainRolled := dice.Roll(RandDrainMax)
-	DrainSize := BaseDrainSize + 16 + 38 + RandDrainRolled
-	readSizeRemain := DrainSize
+
+	drainer, err := drain.NewBehaviorSeedLimitedDrainer(int64(s.userValidator.GetBehaviorSeed()), 16+38, 3266, 64)
+
+	if err != nil {
+		return nil, newError("failed to initialize drainer").Base(err)
+	}
 
 	drainConnection := func(e error) error {
 		// We read a deterministic generated length of data before closing the connection to offset padding read pattern
-		readSizeRemain -= int(buffer.Len())
-		if readSizeRemain > 0 && isDrain {
-			err := s.DrainConnN(reader, readSizeRemain)
-			if err != nil {
-				return newError("failed to drain connection DrainSize = ", BaseDrainSize, " ", RandDrainMax, " ", RandDrainRolled).Base(err).Base(e)
-			}
-			return newError("connection drained DrainSize = ", BaseDrainSize, " ", RandDrainMax, " ", RandDrainRolled).Base(e)
+		drainer.AcknowledgeReceive(int(buffer.Len()))
+		if isDrain {
+			return drain.WithError(drainer, reader, e)
 		}
 		return e
 	}
@@ -182,7 +179,7 @@ func (s *ServerSession) DecodeRequestHeader(reader io.Reader, isDrain bool) (*pr
 		aeadData, shouldDrain, bytesRead, errorReason := vmessaead.OpenVMessAEADHeader(fixedSizeCmdKey, fixedSizeAuthID, reader)
 		if errorReason != nil {
 			if shouldDrain {
-				readSizeRemain -= bytesRead
+				drainer.AcknowledgeReceive(bytesRead)
 				return nil, drainConnection(newError("AEAD read failed").Base(errorReason))
 			} else {
 				return nil, drainConnection(newError("AEAD read failed, drain skipped").Base(errorReason))
@@ -213,7 +210,7 @@ func (s *ServerSession) DecodeRequestHeader(reader io.Reader, isDrain bool) (*pr
 		return nil, drainConnection(newError("invalid user").Base(errorAEAD))
 	}
 
-	readSizeRemain -= int(buffer.Len())
+	drainer.AcknowledgeReceive(int(buffer.Len()))
 	buffer.Clear()
 	if _, err := buffer.ReadFullFrom(decryptor, 38); err != nil {
 		return nil, newError("failed to read request header").Base(err)
@@ -542,8 +539,3 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 		panic("Unknown security type.")
 	}
 }
-
-func (s *ServerSession) DrainConnN(reader io.Reader, n int) error {
-	_, err := io.CopyN(io.Discard, reader, int64(n))
-	return err
-}

proxy/vmess/outbound/outbound.go

@@ -4,6 +4,9 @@ package outbound
 
 import (
 	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"hash/crc64"
 	"time"
 
 	"github.com/xtls/xray-core/transport/internet/stat"
@@ -131,7 +134,12 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 		isAEAD = true
 	}
 
-	session := encoding.NewClientSession(ctx, isAEAD, protocol.DefaultIDHash)
+	hashkdf := hmac.New(sha256.New, []byte("VMessBF"))
+	hashkdf.Write(account.ID.Bytes())
+
+	behaviorSeed := crc64.Checksum(hashkdf.Sum(nil), crc64.MakeTable(crc64.ISO))
+
+	session := encoding.NewClientSession(ctx, isAEAD, protocol.DefaultIDHash, int64(behaviorSeed))
 	sessionPolicy := h.policyManager.ForLevel(request.User.Level)
 
 	ctx, cancel := context.WithCancel(ctx)

testing/scenarios/vmess_test.go

@@ -1437,3 +1437,111 @@ func TestVMessGCMLengthAuth(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestVMessGCMLengthAuthPlusNoTerminationSignal(t *testing.T) {
+	tcpServer := tcp.Server{
+		MsgProcessor: xor,
+	}
+	dest, err := tcpServer.Start()
+	common.Must(err)
+	defer tcpServer.Close()
+
+	userID := protocol.NewID(uuid.New())
+	serverPort := tcp.PickPort()
+	serverConfig := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&log.Config{
+				ErrorLogLevel: clog.Severity_Debug,
+				ErrorLogType:  log.LogType_Console,
+			}),
+		},
+		Inbound: []*core.InboundHandlerConfig{
+			{
+				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
+					PortRange: net.SinglePortRange(serverPort),
+					Listen:    net.NewIPOrDomain(net.LocalHostIP),
+				}),
+				ProxySettings: serial.ToTypedMessage(&inbound.Config{
+					User: []*protocol.User{
+						{
+							Account: serial.ToTypedMessage(&vmess.Account{
+								Id:           userID.String(),
+								AlterId:      64,
+								TestsEnabled: "AuthenticatedLength|NoTerminationSignal",
+							}),
+						},
+					},
+				}),
+			},
+		},
+		Outbound: []*core.OutboundHandlerConfig{
+			{
+				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+			},
+		},
+	}
+
+	clientPort := tcp.PickPort()
+	clientConfig := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&log.Config{
+				ErrorLogLevel: clog.Severity_Debug,
+				ErrorLogType:  log.LogType_Console,
+			}),
+		},
+		Inbound: []*core.InboundHandlerConfig{
+			{
+				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
+					PortRange: net.SinglePortRange(clientPort),
+					Listen:    net.NewIPOrDomain(net.LocalHostIP),
+				}),
+				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
+					Address: net.NewIPOrDomain(dest.Address),
+					Port:    uint32(dest.Port),
+					NetworkList: &net.NetworkList{
+						Network: []net.Network{net.Network_TCP},
+					},
+				}),
+			},
+		},
+		Outbound: []*core.OutboundHandlerConfig{
+			{
+				ProxySettings: serial.ToTypedMessage(&outbound.Config{
+					Receiver: []*protocol.ServerEndpoint{
+						{
+							Address: net.NewIPOrDomain(net.LocalHostIP),
+							Port:    uint32(serverPort),
+							User: []*protocol.User{
+								{
+									Account: serial.ToTypedMessage(&vmess.Account{
+										Id:      userID.String(),
+										AlterId: 64,
+										SecuritySettings: &protocol.SecurityConfig{
+											Type: protocol.SecurityType_AES128_GCM,
+										},
+										TestsEnabled: "AuthenticatedLength|NoTerminationSignal",
+									}),
+								},
+							},
+						},
+					},
+				}),
+			},
+		},
+	}
+
+	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
+	if err != nil {
+		t.Fatal("Failed to initialize all servers: ", err.Error())
+	}
+	defer CloseAllServers(servers)
+
+	var errg errgroup.Group
+	for i := 0; i < 10; i++ {
+		errg.Go(testTCPConn(clientPort, 10240*1024, time.Second*40))
+	}
+
+	if err := errg.Wait(); err != nil {
+		t.Error(err)
+	}
+}