Files
pg/secure/chacha20poly1305/chacha20poly1305.go
2024-09-27 11:20:17 +08:00

104 lines
2.4 KiB
Go

package chacha20poly1305
import (
"crypto/cipher"
"encoding/binary"
"errors"
"sync"
"time"
"golang.org/x/crypto/chacha20poly1305"
"github.com/sigcn/pg/lru"
"github.com/sigcn/pg/secure"
)
var _ secure.SymmAlgo = (*Chacha20Poly1305)(nil)
var timeWindow int64 = 10
func SetDefaultTimeWindow(seconds int64) {
timeWindow = seconds
}
type Chacha20Poly1305 struct {
mut sync.RWMutex
cipher *lru.Cache[string, cipher.AEAD]
provideSecretKey secure.ProvideSecretKey
}
func (s *Chacha20Poly1305) Encrypt(data []byte, pubKey string) ([]byte, error) {
if s == nil {
return nil, errors.New("enc is disabled")
}
aead, err := s.ensureChiperAEAD(pubKey)
if err != nil {
return nil, err
}
nonce := make([]byte, aead.NonceSize())
binary.LittleEndian.PutUint64(nonce[aead.NonceSize()-8:], uint64(time.Now().Unix()/timeWindow))
return aead.Seal(nil, nonce, data, nil), nil
}
func (s *Chacha20Poly1305) Decrypt(data []byte, pubKey string) ([]byte, error) {
if s == nil {
return nil, errors.New("dec is disabled")
}
aead, err := s.ensureChiperAEAD(pubKey)
if err != nil {
return nil, err
}
nonce := make([]byte, aead.NonceSize())
startIndex := aead.NonceSize() - 8
nowUnix := time.Now().Unix()
binary.LittleEndian.PutUint64(nonce[startIndex:], uint64(nowUnix/timeWindow))
plain, err := aead.Open(nil, nonce, data, nil)
if err != nil {
binary.LittleEndian.PutUint64(nonce[startIndex:], uint64(nowUnix/timeWindow+1))
plain, err = aead.Open(nil, nonce, data, nil)
if err != nil {
binary.LittleEndian.PutUint64(nonce[startIndex:], uint64(nowUnix/timeWindow-1))
plain, err = aead.Open(nil, nonce, data, nil)
if err != nil {
return nil, errors.New("invalid data")
}
}
}
return plain, nil
}
func (s *Chacha20Poly1305) SecretKey() secure.ProvideSecretKey {
return s.provideSecretKey
}
func (s *Chacha20Poly1305) ensureChiperAEAD(pubKey string) (cipher.AEAD, error) {
s.mut.RLock()
aead, ok := s.cipher.Get(pubKey)
s.mut.RUnlock()
if !ok {
secretKey, err := s.provideSecretKey(pubKey)
if err != nil {
return nil, err
}
b, err := chacha20poly1305.New(secretKey)
if err != nil {
return nil, err
}
aead = b
s.mut.Lock()
s.cipher.Put(pubKey, aead)
s.mut.Unlock()
}
return aead, nil
}
func New(provideSecretKey secure.ProvideSecretKey) secure.SymmAlgo {
return &Chacha20Poly1305{
cipher: lru.New[string, cipher.AEAD](128),
provideSecretKey: provideSecretKey,
}
}