package sm4

import (
	"crypto/cipher"

	"xdx.jelly/xgcl/gerrors"
)

// Mac 使用SM4CBC加密计算MAC值。即最后一个分组
// 多包的话多次调用即可：data = data1 || data2 || ...
// out = mac(key, iv, data1)
// out = mac(key, out, data2)
// ......
// MAC内不会对输入key, iv, data作任何改动
func Mac(key, iv, data []byte) ([]byte, error) {
	if len(data) < 4096 {
		return mac(key, iv, data)
	}
	return macSmallMem(key, iv, data)
}

// mac make a copy of iv and data, waste a lot of memeory if data is big.
// A little faster.
func mac(key, iv, data []byte) ([]byte, error) {
	dst := make([]byte, len(data)) // dst alloc on stack, for large data
	iv2 := make([]byte, BlockSize) // iv2 escape, alloc on heap
	copy(iv2, iv)

	dst, err := EncryptCBC(dst, iv2, key, data) // assume EncryptCBC not change input iv2
	if err != nil {
		return []byte{}, nil
	}
	copy(iv2, dst[len(dst)-BlockSize:])
	return iv2, nil
}

// macSmallMem for small memeory, local var are only two blocks
func macSmallMem(key, iv, data []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return nil, err
	}
	dst := make([]byte, block.BlockSize()) // dst alloc on stack, for large data
	iv2 := make([]byte, block.BlockSize()) //
	copy(iv2, iv)

	blockMode := cipher.NewCBCEncrypter(block, iv2)
	src := data
	for len(src) >= block.BlockSize() {
		blockMode.CryptBlocks(dst, src[:BlockSize])
		src = src[BlockSize:]
	}
	if len(src) > 0 {
		return nil, gerrors.WithAnnotatingf(ErrSM4MacFailed, "padding first, input length must be multiple of %d", BlockSize)
	}
	return dst, nil

}
