github.com/wfusion/gofusion@v1.1.14/common/utils/cipher/aead_wrapper.go (about)

     1  package cipher
     2  
     3  import (
     4  	"crypto/cipher"
     5  	"errors"
     6  
     7  	"github.com/wfusion/gofusion/common/utils"
     8  )
     9  
    10  const (
    11  	blockSizeTimes   = 1 << 10 // 1024 times
    12  	defaultBlockSize = 16      // 16 bytes
    13  )
    14  
    15  type abstractAEADWrapper struct {
    16  	cipher.AEAD
    17  	overheadSize    int
    18  	nonceSize       int
    19  	blockSize       int
    20  	plainBlockSize  int
    21  	cipherBlockSize int
    22  	sealedSize      int
    23  
    24  	mode Mode
    25  }
    26  
    27  func newAbstractAEADWrapper(cipherAEAD cipher.AEAD, blockSize int, mode Mode) abstractAEADWrapper {
    28  	plainBlockSize := blockSize * blockSizeTimes
    29  	return abstractAEADWrapper{
    30  		AEAD:            cipherAEAD,
    31  		overheadSize:    cipherAEAD.Overhead(),
    32  		nonceSize:       cipherAEAD.NonceSize(),
    33  		blockSize:       blockSize,
    34  		plainBlockSize:  plainBlockSize,
    35  		cipherBlockSize: cipherAEAD.NonceSize() + cipherAEAD.Overhead() + plainBlockSize,
    36  		sealedSize:      cipherAEAD.Overhead() + plainBlockSize,
    37  		mode:            mode,
    38  	}
    39  }
    40  
    41  func (a *abstractAEADWrapper) PlainBlockSize() int  { return a.plainBlockSize }
    42  func (a *abstractAEADWrapper) CipherBlockSize() int { return a.cipherBlockSize }
    43  func (a *abstractAEADWrapper) CipherMode() Mode     { return a.mode }
    44  
    45  type aeadEncryptWrapper struct {
    46  	abstractAEADWrapper
    47  }
    48  
    49  func newAEADEncryptWrapper(cipherAEAD cipher.AEAD, blockSize int, mode Mode) blockMode {
    50  	return &aeadEncryptWrapper{abstractAEADWrapper: newAbstractAEADWrapper(cipherAEAD, blockSize, mode)}
    51  }
    52  
    53  func (a *aeadEncryptWrapper) CryptBlocks(dst, src, buf []byte) (n int, err error) {
    54  	if len(src) == 0 {
    55  		return
    56  	}
    57  
    58  	nonce, cb := utils.BytesPool.Get(a.nonceSize)
    59  	defer cb()
    60  	if _, err = utils.CryptoRandom(nonce); err != nil {
    61  		return
    62  	}
    63  
    64  	sealed := a.Seal(buf[:0], nonce, src[:utils.Min(a.plainBlockSize, len(src))], nil)
    65  	copy(dst[:a.nonceSize], nonce)
    66  	copy(dst[a.nonceSize:], sealed)
    67  	n = len(nonce) + len(sealed)
    68  	return
    69  }
    70  
    71  type aeadDecryptWrapper struct {
    72  	abstractAEADWrapper
    73  }
    74  
    75  func newAEADDecryptWrapper(cipherAEAD cipher.AEAD, blockSize int, mode Mode) blockMode {
    76  	return &aeadDecryptWrapper{abstractAEADWrapper: newAbstractAEADWrapper(cipherAEAD, blockSize, mode)}
    77  }
    78  
    79  func (a *aeadDecryptWrapper) CryptBlocks(dst, src, buf []byte) (n int, err error) {
    80  	if len(src) == 0 {
    81  		return
    82  	}
    83  	sealedSize := utils.Min(a.sealedSize, len(src)-a.nonceSize)
    84  	if sealedSize < 0 {
    85  		return 0, errors.New("input not full blocks when decrypt")
    86  	}
    87  
    88  	nonce := src[:a.nonceSize]
    89  	src = src[a.nonceSize:]
    90  	unsealed, err := a.Open(buf[:0], nonce, src[:sealedSize], nil)
    91  	if err != nil {
    92  		return
    93  	}
    94  
    95  	copy(dst, unsealed)
    96  	n = len(unsealed)
    97  	return
    98  }