github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/shadowsocks/core/cipher.go (about)

     1  package core
     2  
     3  import (
     4  	"crypto/md5"
     5  	"errors"
     6  	"net"
     7  	"sort"
     8  	"strings"
     9  
    10  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/shadowsocks/shadowaead"
    11  )
    12  
    13  type Cipher interface {
    14  	StreamConnCipher
    15  	PacketConnCipher
    16  }
    17  
    18  type StreamConnCipher interface {
    19  	StreamConn(net.Conn) net.Conn
    20  }
    21  
    22  type PacketConnCipher interface {
    23  	PacketConn(net.PacketConn) net.PacketConn
    24  }
    25  
    26  // ErrCipherNotSupported occurs when a cipher is not supported (likely because of security concerns).
    27  var ErrCipherNotSupported = errors.New("cipher not supported")
    28  
    29  const (
    30  	aeadAes128Gcm        = "AEAD_AES_128_GCM"
    31  	aeadAes256Gcm        = "AEAD_AES_256_GCM"
    32  	aeadChacha20Poly1305 = "AEAD_CHACHA20_POLY1305"
    33  )
    34  
    35  // List of AEAD ciphers: key size in bytes and constructor
    36  var aeadList = map[string]struct {
    37  	KeySize int
    38  	New     func([]byte) (shadowaead.Cipher, error)
    39  }{
    40  	aeadAes128Gcm:        {16, shadowaead.AESGCM},
    41  	aeadAes256Gcm:        {32, shadowaead.AESGCM},
    42  	aeadChacha20Poly1305: {32, shadowaead.Chacha20Poly1305},
    43  }
    44  
    45  // ListCipher returns a list of available cipher names sorted alphabetically.
    46  func ListCipher() []string {
    47  	var l []string
    48  	for k := range aeadList {
    49  		l = append(l, k)
    50  	}
    51  	sort.Strings(l)
    52  	return l
    53  }
    54  
    55  // PickCipher returns a Cipher of the given name. Derive key from password if given key is empty.
    56  func PickCipher(name string, key []byte, password string) (Cipher, error) {
    57  	name = strings.ToUpper(name)
    58  
    59  	switch name {
    60  	case "DUMMY":
    61  		return &dummy{}, nil
    62  	case "CHACHA20-IETF-POLY1305":
    63  		name = aeadChacha20Poly1305
    64  	case "AES-128-GCM":
    65  		name = aeadAes128Gcm
    66  	case "AES-256-GCM":
    67  		name = aeadAes256Gcm
    68  	}
    69  
    70  	if choice, ok := aeadList[name]; ok {
    71  		if len(key) == 0 {
    72  			key = KDF(password, choice.KeySize)
    73  		}
    74  		if len(key) != choice.KeySize {
    75  			return nil, shadowaead.KeySizeError(choice.KeySize)
    76  		}
    77  		aead, err := choice.New(key)
    78  		return &aeadCipher{aead}, err
    79  	}
    80  
    81  	return nil, ErrCipherNotSupported
    82  }
    83  
    84  type aeadCipher struct{ shadowaead.Cipher }
    85  
    86  func (aead *aeadCipher) StreamConn(c net.Conn) net.Conn { return shadowaead.NewConn(c, aead) }
    87  func (aead *aeadCipher) PacketConn(c net.PacketConn) net.PacketConn {
    88  	return shadowaead.NewPacketConn(c, aead)
    89  }
    90  
    91  // dummy cipher does not encrypt
    92  type dummy struct{}
    93  
    94  func (dummy) StreamConn(c net.Conn) net.Conn             { return c }
    95  func (dummy) PacketConn(c net.PacketConn) net.PacketConn { return c }
    96  
    97  // key-derivation function from original Shadowsocks
    98  func KDF(password string, keyLen int) []byte {
    99  	var b, prev []byte
   100  	h := md5.New()
   101  	for len(b) < keyLen {
   102  		h.Write(prev)
   103  		h.Write([]byte(password))
   104  		b = h.Sum(b)
   105  		prev = b[len(b)-h.Size():]
   106  		h.Reset()
   107  	}
   108  	return b[:keyLen]
   109  }