Feature: support snell v2 (#952)

Co-authored-by: Dreamacro <8615343+Dreamacro@users.noreply.github.com>

component/snell/cipher.go

@@ -1,21 +1,54 @@
 package snell
 
 import (
+	"crypto/aes"
 	"crypto/cipher"
 
+	"github.com/Dreamacro/go-shadowsocks2/shadowaead"
 	"golang.org/x/crypto/argon2"
+	"golang.org/x/crypto/chacha20poly1305"
 )
 
 type snellCipher struct {
 	psk      []byte
+	keySize  int
 	makeAEAD func(key []byte) (cipher.AEAD, error)
 }
 
-func (sc *snellCipher) KeySize() int  { return 32 }
+func (sc *snellCipher) KeySize() int  { return sc.keySize }
 func (sc *snellCipher) SaltSize() int { return 16 }
 func (sc *snellCipher) Encrypter(salt []byte) (cipher.AEAD, error) {
-	return sc.makeAEAD(argon2.IDKey(sc.psk, salt, 3, 8, 1, uint32(sc.KeySize())))
+	return sc.makeAEAD(snellKDF(sc.psk, salt, sc.KeySize()))
 }
 func (sc *snellCipher) Decrypter(salt []byte) (cipher.AEAD, error) {
-	return sc.makeAEAD(argon2.IDKey(sc.psk, salt, 3, 8, 1, uint32(sc.KeySize())))
+	return sc.makeAEAD(snellKDF(sc.psk, salt, sc.KeySize()))
+}
+
+func snellKDF(psk, salt []byte, keySize int) []byte {
+	// snell use a special kdf function
+	return argon2.IDKey(psk, salt, 3, 8, 1, 32)[:keySize]
+}
+
+func aesGCM(key []byte) (cipher.AEAD, error) {
+	blk, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+	return cipher.NewGCM(blk)
+}
+
+func NewAES128GCM(psk []byte) shadowaead.Cipher {
+	return &snellCipher{
+		psk:      psk,
+		keySize:  16,
+		makeAEAD: aesGCM,
+	}
+}
+
+func NewChacha20Poly1305(psk []byte) shadowaead.Cipher {
+	return &snellCipher{
+		psk:      psk,
+		keySize:  32,
+		makeAEAD: chacha20poly1305.New,
+	}
 }

component/snell/pool.go

@@ -0,0 +1,80 @@
+package snell
+
+import (
+	"context"
+	"net"
+
+	"github.com/Dreamacro/clash/component/pool"
+	"github.com/Dreamacro/go-shadowsocks2/shadowaead"
+)
+
+type Pool struct {
+	pool *pool.Pool
+}
+
+func (p *Pool) Get() (net.Conn, error) {
+	return p.GetContext(context.Background())
+}
+
+func (p *Pool) GetContext(ctx context.Context) (net.Conn, error) {
+	elm, err := p.pool.GetContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return &PoolConn{elm.(*Snell), p}, nil
+}
+
+func (p *Pool) Put(conn net.Conn) {
+	if err := HalfClose(conn); err != nil {
+		conn.Close()
+		return
+	}
+
+	p.pool.Put(conn)
+}
+
+type PoolConn struct {
+	*Snell
+	pool *Pool
+}
+
+func (pc *PoolConn) Read(b []byte) (int, error) {
+	// save old status of reply (it mutable by Read)
+	reply := pc.Snell.reply
+
+	n, err := pc.Snell.Read(b)
+	if err == shadowaead.ErrZeroChunk {
+		// if reply is false, it should be client halfclose.
+		// ignore error and read data again.
+		if !reply {
+			pc.Snell.reply = false
+			return pc.Snell.Read(b)
+		}
+	}
+	return n, err
+}
+
+func (pc *PoolConn) Write(b []byte) (int, error) {
+	return pc.Snell.Write(b)
+}
+
+func (pc *PoolConn) Close() error {
+	pc.pool.Put(pc.Snell)
+	return nil
+}
+
+func NewPool(factory func(context.Context) (*Snell, error)) *Pool {
+	p := pool.New(
+		func(ctx context.Context) (interface{}, error) {
+			return factory(ctx)
+		},
+		pool.WithAge(15000),
+		pool.WithSize(10),
+		pool.WithEvict(func(item interface{}) {
+			item.(*Snell).Close()
+		}),
+	)
+
+	return &Pool{p}
+}

component/snell/snell.go

@@ -10,14 +10,21 @@ import (
 	"sync"
 
 	"github.com/Dreamacro/go-shadowsocks2/shadowaead"
-	"golang.org/x/crypto/chacha20poly1305"
 )
 
 const (
-	CommandPing    byte = 0
-	CommandConnect byte = 1
+	Version1            = 1
+	Version2            = 2
+	DefaultSnellVersion = Version1
+)
+
+const (
+	CommandPing      byte = 0
+	CommandConnect   byte = 1
+	CommandConnectV2 byte = 5
 
 	CommandTunnel byte = 0
+	CommandPong   byte = 1
 	CommandError  byte = 2
 
 	Version byte = 1
@@ -25,6 +32,7 @@ const (
 
 var (
 	bufferPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }}
+	endSignal  = []byte{}
 )
 
 type Snell struct {
@@ -70,12 +78,16 @@ func (s *Snell) Read(b []byte) (int, error) {
 	return 0, fmt.Errorf("server reported code: %d, message: %s", errcode, string(msg))
 }
 
-func WriteHeader(conn net.Conn, host string, port uint) error {
+func WriteHeader(conn net.Conn, host string, port uint, version int) error {
 	buf := bufferPool.Get().(*bytes.Buffer)
 	buf.Reset()
 	defer bufferPool.Put(buf)
 	buf.WriteByte(Version)
-	buf.WriteByte(CommandConnect)
+	if version == Version2 {
+		buf.WriteByte(CommandConnectV2)
+	} else {
+		buf.WriteByte(CommandConnect)
+	}
 
 	// clientID length & id
 	buf.WriteByte(0)
@@ -92,7 +104,24 @@ func WriteHeader(conn net.Conn, host string, port uint) error {
 	return nil
 }
 
-func StreamConn(conn net.Conn, psk []byte) net.Conn {
-	cipher := &snellCipher{psk, chacha20poly1305.New}
+// HalfClose works only on version2
+func HalfClose(conn net.Conn) error {
+	if _, err := conn.Write(endSignal); err != nil {
+		return err
+	}
+
+	if s, ok := conn.(*Snell); ok {
+		s.reply = false
+	}
+	return nil
+}
+
+func StreamConn(conn net.Conn, psk []byte, version int) *Snell {
+	var cipher shadowaead.Cipher
+	if version == Version2 {
+		cipher = NewAES128GCM(psk)
+	} else {
+		cipher = NewChacha20Poly1305(psk)
+	}
 	return &Snell{Conn: shadowaead.NewConn(conn, cipher)}
 }