Feature: support snell v2 (#952)

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

component/pool/pool.go

@@ -0,0 +1,114 @@
+package pool
+
+import (
+	"context"
+	"runtime"
+	"time"
+)
+
+type Factory = func(context.Context) (interface{}, error)
+
+type entry struct {
+	elm  interface{}
+	time time.Time
+}
+
+type Option func(*pool)
+
+// WithEvict set the evict callback
+func WithEvict(cb func(interface{})) Option {
+	return func(p *pool) {
+		p.evict = cb
+	}
+}
+
+// WithAge defined element max age (millisecond)
+func WithAge(maxAge int64) Option {
+	return func(p *pool) {
+		p.maxAge = maxAge
+	}
+}
+
+// WithSize defined max size of Pool
+func WithSize(maxSize int) Option {
+	return func(p *pool) {
+		p.ch = make(chan interface{}, maxSize)
+	}
+}
+
+// Pool is for GC, see New for detail
+type Pool struct {
+	*pool
+}
+
+type pool struct {
+	ch      chan interface{}
+	factory Factory
+	evict   func(interface{})
+	maxAge  int64
+}
+
+func (p *pool) GetContext(ctx context.Context) (interface{}, error) {
+	now := time.Now()
+	for {
+		select {
+		case item := <-p.ch:
+			elm := item.(*entry)
+			if p.maxAge != 0 && now.Sub(item.(*entry).time).Milliseconds() > p.maxAge {
+				if p.evict != nil {
+					p.evict(elm.elm)
+				}
+				continue
+			}
+
+			return elm.elm, nil
+		default:
+			return p.factory(ctx)
+		}
+	}
+}
+
+func (p *pool) Get() (interface{}, error) {
+	return p.GetContext(context.Background())
+}
+
+func (p *pool) Put(item interface{}) {
+	e := &entry{
+		elm:  item,
+		time: time.Now(),
+	}
+
+	select {
+	case p.ch <- e:
+		return
+	default:
+		// pool is full
+		if p.evict != nil {
+			p.evict(item)
+		}
+		return
+	}
+}
+
+func recycle(p *Pool) {
+	for item := range p.pool.ch {
+		if p.pool.evict != nil {
+			p.pool.evict(item.(*entry).elm)
+		}
+	}
+}
+
+func New(factory Factory, options ...Option) *Pool {
+	p := &pool{
+		ch:      make(chan interface{}, 10),
+		factory: factory,
+	}
+
+	for _, option := range options {
+		option(p)
+	}
+
+	P := &Pool{p}
+	runtime.SetFinalizer(P, recycle)
+	return P
+}

component/pool/pool_test.go

@@ -0,0 +1,96 @@
+package pool
+
+import (
+	"context"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func lg() Factory {
+	initial := -1
+	return func(context.Context) (interface{}, error) {
+		initial++
+		return initial, nil
+	}
+}
+
+func TestPool_Basic(t *testing.T) {
+	g := lg()
+	pool := New(g)
+
+	elm, _ := pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(elm)
+	elm, _ = pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	elm, _ = pool.Get()
+	assert.Equal(t, 1, elm.(int))
+}
+
+func TestPool_MaxSize(t *testing.T) {
+	g := lg()
+	size := 5
+	pool := New(g, WithSize(size))
+
+	items := []interface{}{}
+
+	for i := 0; i < size; i++ {
+		item, _ := pool.Get()
+		items = append(items, item)
+	}
+
+	extra, _ := pool.Get()
+	assert.Equal(t, size, extra.(int))
+
+	for _, item := range items {
+		pool.Put(item)
+	}
+
+	pool.Put(extra)
+
+	for _, item := range items {
+		elm, _ := pool.Get()
+		assert.Equal(t, item.(int), elm.(int))
+	}
+}
+
+func TestPool_MaxAge(t *testing.T) {
+	g := lg()
+	pool := New(g, WithAge(20))
+
+	elm, _ := pool.Get()
+	pool.Put(elm)
+
+	elm, _ = pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(elm)
+
+	time.Sleep(time.Millisecond * 22)
+	elm, _ = pool.Get()
+	assert.Equal(t, 1, elm.(int))
+}
+
+func TestPool_AutoGC(t *testing.T) {
+	g := lg()
+
+	sign := make(chan int)
+	pool := New(g, WithEvict(func(item interface{}) {
+		sign <- item.(int)
+	}))
+
+	elm, _ := pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(2)
+
+	runtime.GC()
+
+	select {
+	case num := <-sign:
+		assert.Equal(t, 2, num)
+	case <-time.After(time.Second * 3):
+		assert.Fail(t, "something wrong")
+	}
+}

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)}
 }