github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/shadowsocksr/utils/tool.go (about)

     1  package ssr
     2  
     3  import (
     4  	"crypto"
     5  	"hash"
     6  	"sync"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/utils/syncmap"
     9  )
    10  
    11  type HMAC crypto.Hash
    12  
    13  func (h HMAC) HMAC(key, data, buf []byte) []byte { return Hmac(crypto.Hash(h), key, data, buf) }
    14  func (h HMAC) HASH(b []byte) []byte              { return HashSum(crypto.Hash(h), b) }
    15  
    16  var hmacPool syncmap.SyncMap[crypto.Hash, *sync.Pool]
    17  
    18  func getHmac(hash crypto.Hash, key []byte) CHmac {
    19  	z, ok := hmacPool.Load(hash)
    20  	if !ok {
    21  		z = &sync.Pool{New: func() any { return NewHmac(hash) }}
    22  		hmacPool.Store(hash, z)
    23  	}
    24  
    25  	h := z.Get().(CHmac)
    26  	h.ResetKey(key)
    27  	return h
    28  }
    29  
    30  func putHmac(h CHmac) {
    31  	z, ok := hmacPool.Load(h.Hash())
    32  	if !ok {
    33  
    34  		z = &sync.Pool{New: func() any { return NewHmac(h.Hash()) }}
    35  		hmacPool.Store(h.Hash(), z)
    36  	}
    37  	z.Put(h)
    38  }
    39  
    40  func Hmac(c crypto.Hash, key, data, buf []byte) []byte {
    41  	h := getHmac(c, key)
    42  	defer putHmac(h)
    43  
    44  	h.Write(data)
    45  
    46  	if buf == nil {
    47  		buf = make([]byte, h.Size())
    48  	}
    49  
    50  	copy(buf, h.Sum(nil))
    51  	return buf
    52  }
    53  
    54  var hashPool syncmap.SyncMap[crypto.Hash, *sync.Pool]
    55  
    56  type chash struct {
    57  	hash.Hash
    58  	h crypto.Hash
    59  }
    60  
    61  func (c chash) CryptoHash() crypto.Hash { return c.h }
    62  
    63  func newCHash(c crypto.Hash) *chash { return &chash{h: c, Hash: c.New()} }
    64  
    65  func getHash(ha crypto.Hash) *chash {
    66  	h, ok := hashPool.Load(ha)
    67  	if !ok {
    68  		h = &sync.Pool{New: func() any { return newCHash(ha) }}
    69  		hashPool.Store(ha, h)
    70  	}
    71  	z := h.Get().(*chash)
    72  	z.Reset()
    73  	return z
    74  }
    75  
    76  func putHash(hh *chash) {
    77  	h, ok := hashPool.Load(hh.CryptoHash())
    78  	if !ok {
    79  		h = &sync.Pool{New: func() any { return newCHash(hh.CryptoHash()) }}
    80  		hashPool.Store(hh.CryptoHash(), h)
    81  	}
    82  	hh.Reset()
    83  	h.Put(hh)
    84  }
    85  
    86  func HashSum(h crypto.Hash, d []byte) []byte {
    87  	hh := getHash(h)
    88  	defer putHash(hh)
    89  	hh.Reset()
    90  	hh.Write(d)
    91  	return hh.Sum(nil)
    92  }