fix core

core/cipher.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/BarryDeng/go-shadowsocks2/shadowaead"
+	"github.com/BarryDeng/go-shadowsocks2/shadowstream"
 )
 
 type Cipher interface {
@@ -27,19 +28,41 @@ type PacketConnCipher interface {
 var ErrCipherNotSupported = errors.New("cipher not supported")
 
 const (
-	aeadAes128Gcm        = "AEAD_AES_128_GCM"
-	aeadAes256Gcm        = "AEAD_AES_256_GCM"
-	aeadChacha20Poly1305 = "AEAD_CHACHA20_POLY1305"
+	aeadAes128Gcm         = "AEAD_AES_128_GCM"
+	aeadAes192Gcm         = "AEAD_AES_192_GCM"
+	aeadAes256Gcm         = "AEAD_AES_256_GCM"
+	aeadChacha20Poly1305  = "AEAD_CHACHA20_POLY1305"
+	aeadXChacha20Poly1305 = "AEAD_XCHACHA20_POLY1305"
+	sm4128Gcm             = "SM4_128_GCM"
 )
 
 // List of AEAD ciphers: key size in bytes and constructor
 var aeadList = map[string]struct {
 	KeySize int
 	New     func([]byte) (shadowaead.Cipher, error)
 }{
-	aeadAes128Gcm:        {16, shadowaead.AESGCM},
-	aeadAes256Gcm:        {32, shadowaead.AESGCM},
-	aeadChacha20Poly1305: {32, shadowaead.Chacha20Poly1305},
+	aeadAes128Gcm:         {16, shadowaead.AESGCM},
+	aeadAes192Gcm:         {24, shadowaead.AESGCM},
+	aeadAes256Gcm:         {32, shadowaead.AESGCM},
+	aeadChacha20Poly1305:  {32, shadowaead.Chacha20Poly1305},
+	aeadXChacha20Poly1305: {32, shadowaead.XChacha20Poly1305},
+	sm4128Gcm:             {16, shadowaead.SM4GCM},
+}
+
+// List of stream ciphers: key size in bytes and constructor
+var streamList = map[string]struct {
+	KeySize int
+	New     func(key []byte) (shadowstream.Cipher, error)
+}{
+	"RC4-MD5":       {16, shadowstream.RC4MD5},
+	"AES-128-CTR":   {16, shadowstream.AESCTR},
+	"AES-192-CTR":   {24, shadowstream.AESCTR},
+	"AES-256-CTR":   {32, shadowstream.AESCTR},
+	"AES-128-CFB":   {16, shadowstream.AESCFB},
+	"AES-192-CFB":   {24, shadowstream.AESCFB},
+	"AES-256-CFB":   {32, shadowstream.AESCFB},
+	"CHACHA20-IETF": {32, shadowstream.Chacha20IETF},
+	"XCHACHA20":     {32, shadowstream.Xchacha20},
 }
 
 // ListCipher returns a list of available cipher names sorted alphabetically.
@@ -61,41 +84,73 @@ func PickCipher(name string, key []byte, password string) (Cipher, error) {
 		return &dummy{}, nil
 	case "CHACHA20-IETF-POLY1305":
 		name = aeadChacha20Poly1305
+	case "XCHACHA20-IETF-POLY1305":
+		name = aeadXChacha20Poly1305
 	case "AES-128-GCM":
 		name = aeadAes128Gcm
+	case "AES-192-GCM":
+		name = aeadAes192Gcm
 	case "AES-256-GCM":
 		name = aeadAes256Gcm
+	case "SM4-128-GCM":
+		name = sm4128Gcm
 	}
 
 	if choice, ok := aeadList[name]; ok {
 		if len(key) == 0 {
-			key = kdf(password, choice.KeySize)
+			key = Kdf(password, choice.KeySize)
 		}
 		if len(key) != choice.KeySize {
 			return nil, shadowaead.KeySizeError(choice.KeySize)
 		}
 		aead, err := choice.New(key)
-		return &aeadCipher{aead}, err
+		return &AeadCipher{Cipher: aead, Key: key}, err
+	}
+
+	if choice, ok := streamList[name]; ok {
+		if len(key) == 0 {
+			key = Kdf(password, choice.KeySize)
+		}
+		if len(key) != choice.KeySize {
+			return nil, shadowstream.KeySizeError(choice.KeySize)
+		}
+		ciph, err := choice.New(key)
+		return &StreamCipher{Cipher: ciph, Key: key}, err
 	}
 
 	return nil, ErrCipherNotSupported
 }
 
-type aeadCipher struct{ shadowaead.Cipher }
+type AeadCipher struct {
+	shadowaead.Cipher
+
+	Key []byte
+}
 
-func (aead *aeadCipher) StreamConn(c net.Conn) net.Conn { return shadowaead.NewConn(c, aead) }
-func (aead *aeadCipher) PacketConn(c net.PacketConn) net.PacketConn {
+func (aead *AeadCipher) StreamConn(c net.Conn) net.Conn { return shadowaead.NewConn(c, aead) }
+func (aead *AeadCipher) PacketConn(c net.PacketConn) net.PacketConn {
 	return shadowaead.NewPacketConn(c, aead)
 }
 
+type StreamCipher struct {
+	shadowstream.Cipher
+
+	Key []byte
+}
+
+func (ciph *StreamCipher) StreamConn(c net.Conn) net.Conn { return shadowstream.NewConn(c, ciph) }
+func (ciph *StreamCipher) PacketConn(c net.PacketConn) net.PacketConn {
+	return shadowstream.NewPacketConn(c, ciph)
+}
+
 // dummy cipher does not encrypt
 type dummy struct{}
 
 func (dummy) StreamConn(c net.Conn) net.Conn             { return c }
 func (dummy) PacketConn(c net.PacketConn) net.PacketConn { return c }
 
 // key-derivation function from original Shadowsocks
-func kdf(password string, keyLen int) []byte {
+func Kdf(password string, keyLen int) []byte {
 	var b, prev []byte
 	h := md5.New()
 	for len(b) < keyLen {