Chore: split component to transport

transport/snell/cipher.go

@@ -0,0 +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 sc.keySize }
+func (sc *snellCipher) SaltSize() int { return 16 }
+func (sc *snellCipher) Encrypter(salt []byte) (cipher.AEAD, error) {
+	return sc.makeAEAD(snellKDF(sc.psk, salt, sc.KeySize()))
+}
+func (sc *snellCipher) Decrypter(salt []byte) (cipher.AEAD, error) {
+	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,
+	}
+}

transport/snell/pool.go

@@ -0,0 +1,85 @@
+package snell
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"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 {
+	// clash use SetReadDeadline to break bidirectional copy between client and server.
+	// reset it before reuse connection to avoid io timeout error.
+	pc.Snell.Conn.SetReadDeadline(time.Time{})
+	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}
+}

transport/snell/snell.go

@@ -0,0 +1,127 @@
+package snell
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+
+	"github.com/Dreamacro/go-shadowsocks2/shadowaead"
+)
+
+const (
+	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
+)
+
+var (
+	bufferPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }}
+	endSignal  = []byte{}
+)
+
+type Snell struct {
+	net.Conn
+	buffer [1]byte
+	reply  bool
+}
+
+func (s *Snell) Read(b []byte) (int, error) {
+	if s.reply {
+		return s.Conn.Read(b)
+	}
+
+	s.reply = true
+	if _, err := io.ReadFull(s.Conn, s.buffer[:]); err != nil {
+		return 0, err
+	}
+
+	if s.buffer[0] == CommandTunnel {
+		return s.Conn.Read(b)
+	} else if s.buffer[0] != CommandError {
+		return 0, errors.New("command not support")
+	}
+
+	// CommandError
+	// 1 byte error code
+	if _, err := io.ReadFull(s.Conn, s.buffer[:]); err != nil {
+		return 0, err
+	}
+	errcode := int(s.buffer[0])
+
+	// 1 byte error message length
+	if _, err := io.ReadFull(s.Conn, s.buffer[:]); err != nil {
+		return 0, err
+	}
+	length := int(s.buffer[0])
+	msg := make([]byte, length)
+
+	if _, err := io.ReadFull(s.Conn, msg); err != nil {
+		return 0, err
+	}
+
+	return 0, fmt.Errorf("server reported code: %d, message: %s", errcode, string(msg))
+}
+
+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)
+	if version == Version2 {
+		buf.WriteByte(CommandConnectV2)
+	} else {
+		buf.WriteByte(CommandConnect)
+	}
+
+	// clientID length & id
+	buf.WriteByte(0)
+
+	// host & port
+	buf.WriteByte(uint8(len(host)))
+	buf.WriteString(host)
+	binary.Write(buf, binary.BigEndian, uint16(port))
+
+	if _, err := conn.Write(buf.Bytes()); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// 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)}
+}