mirror of
https://github.com/XTLS/Xray-core.git
synced 2025-04-29 16:58:34 +00:00
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>
This commit is contained in:
parent
6fb5c887b2
commit
d5a7901601
9 changed files with 277 additions and 66 deletions
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue