added drain for ss and improved vmess drain

proxy/shadowsocks/protocol.go

@@ -4,8 +4,14 @@ package shadowsocks
 
 import (
 	"bytes"
+	"crypto/hmac"
 	"crypto/rand"
+	"crypto/sha256"
+	"hash"
+	"hash/crc32"
 	"io"
+	"io/ioutil"
+	"v2ray.com/core/common/dice"
 
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/bitmask"
@@ -32,6 +38,19 @@ var addrParser = protocol.NewAddressParser(
 func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.RequestHeader, buf.Reader, error) {
 	account := user.Account.(*MemoryAccount)
 
+	hashkdf := hmac.New(func()hash.Hash{return sha256.New()}, []byte("SSBSKDF"))
+	hashkdf.Write(account.Key)
+
+	behaviorSeed := crc32.ChecksumIEEE(hashkdf.Sum(nil))
+
+	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
+
+
 	buffer := buf.New()
 	defer buffer.Release()
 
@@ -39,6 +58,8 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 	var iv []byte
 	if ivLen > 0 {
 		if _, err := buffer.ReadFullFrom(reader, ivLen); err != nil {
+			readSizeRemain -= int(buffer.Len())
+			DrainConnN(reader,readSizeRemain)
 			return nil, nil, newError("failed to read IV").Base(err)
 		}
 
@@ -47,6 +68,8 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 
 	r, err := account.Cipher.NewDecryptionReader(account.Key, iv, reader)
 	if err != nil {
+		readSizeRemain -= int(buffer.Len())
+		DrainConnN(reader,readSizeRemain)
 		return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
 	}
 	br := &buf.BufferedReader{Reader: r}
@@ -58,10 +81,13 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 		Command: protocol.RequestCommandTCP,
 	}
 
+	readSizeRemain -= int(buffer.Len())
 	buffer.Clear()
 
 	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)
 	}
 
@@ -74,10 +100,14 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 		}
 
 		if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
+			readSizeRemain -= int(buffer.Len())
+			DrainConnN(reader,readSizeRemain)
 			return nil, nil, newError("rejecting connection with OTA enabled, while server disables OTA")
 		}
 
 		if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
+			readSizeRemain -= int(buffer.Len())
+			DrainConnN(reader,readSizeRemain)
 			return nil, nil, newError("rejecting connection with OTA disabled, while server enables OTA")
 		}
 	}
@@ -88,15 +118,21 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 
 		_, err := buffer.ReadFullFrom(br, AuthSize)
 		if err != nil {
+			readSizeRemain -= int(buffer.Len())
+			DrainConnN(reader,readSizeRemain)
 			return nil, nil, newError("Failed to read OTA").Base(err)
 		}
 
 		if !bytes.Equal(actualAuth, buffer.BytesFrom(-AuthSize)) {
+			readSizeRemain -= int(buffer.Len())
+			DrainConnN(reader,readSizeRemain)
 			return nil, nil, newError("invalid OTA")
 		}
 	}
 
 	if request.Address == nil {
+		readSizeRemain -= int(buffer.Len())
+		DrainConnN(reader,readSizeRemain)
 		return nil, nil, newError("invalid remote address.")
 	}
 
@@ -110,6 +146,11 @@ func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.Requ
 	return request, chunkReader, nil
 }
 
+func DrainConnN(reader io.Reader, n int) error {
+	_, err := io.CopyN(ioutil.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

proxy/vmess/validator.go

@@ -3,6 +3,9 @@
 package vmess
 
 import (
+	"crypto/hmac"
+	"crypto/sha256"
+	"hash"
 	"hash/crc64"
 	"strings"
 	"sync"
@@ -139,7 +142,9 @@ func (v *TimedUserValidator) Add(u *protocol.MemoryUser) error {
 
 	account := uu.user.Account.(*MemoryAccount)
 	if v.behaviorFused == false {
-		v.behaviorSeed = crc64.Update(v.behaviorSeed, crc64.MakeTable(crc64.ECMA), account.ID.Bytes())
+		hashkdf := hmac.New(func()hash.Hash{return sha256.New()}, []byte("VMESSBSKDF"))
+		hashkdf.Write(account.ID.Bytes())
+		v.behaviorSeed = crc64.Update(v.behaviorSeed, crc64.MakeTable(crc64.ECMA), hashkdf.Sum(nil))
 	}
 
 	var cmdkeyfl [16]byte