Skip to content

Commit

Permalink
Improve the performance of AES-CTR-HMAC by avoiding unnecessary alloc…
Browse files Browse the repository at this point in the history
…ations and copies.

PiperOrigin-RevId: 705911944
Change-Id: I3a6a2deb8ae12f89454e95bbb9526f10fdce74da
  • Loading branch information
fernandolobato authored and copybara-github committed Dec 13, 2024
1 parent b87aa94 commit e18b6d1
Show file tree
Hide file tree
Showing 9 changed files with 783 additions and 168 deletions.
74 changes: 55 additions & 19 deletions aead/aesctrhmac/aead.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,62 +16,98 @@ package aesctrhmac

import (
"bytes"
"encoding/binary"
"errors"
"fmt"
"slices"

"github.com/tink-crypto/tink-go/v2/aead/subtle"
"github.com/tink-crypto/tink-go/v2/insecuresecretdataaccess"
"github.com/tink-crypto/tink-go/v2/internal/aead"
"github.com/tink-crypto/tink-go/v2/internal/mac/hmac"
"github.com/tink-crypto/tink-go/v2/key"
subtlemac "github.com/tink-crypto/tink-go/v2/mac/subtle"
"github.com/tink-crypto/tink-go/v2/tink"
)

type fullAEAD struct {
aead *subtle.EncryptThenAuthenticate
aesCTR *aead.AESCTR
hmac *hmac.HMAC
ivSize int
tagSize int
prefix []byte
variant Variant
}

var _ tink.AEAD = (*fullAEAD)(nil)

func newAEAD(key *Key) (tink.AEAD, error) {
aesCTR, err := subtle.NewAESCTR(key.AESKeyBytes().Data(insecuresecretdataaccess.Token{}), key.parameters.IVSizeInBytes())
tagSize := key.parameters.TagSizeInBytes()
ivSize := key.parameters.IVSizeInBytes()
aesCTR, err := aead.NewAESCTR(key.AESKeyBytes().Data(insecuresecretdataaccess.Token{}), ivSize)
if err != nil {
return nil, err
}
hmac, err := subtlemac.NewHMAC(key.parameters.HashType().String(), key.HMACKeyBytes().Data(insecuresecretdataaccess.Token{}), uint32(key.parameters.TagSizeInBytes()))
if err != nil {
return nil, err
}
eta, err := subtle.NewEncryptThenAuthenticate(aesCTR, hmac, key.parameters.TagSizeInBytes())
hmac, err := hmac.New(key.parameters.HashType().String(), key.HMACKeyBytes().Data(insecuresecretdataaccess.Token{}), uint32(tagSize))
if err != nil {
return nil, err
}
return &fullAEAD{
aead: eta,
aesCTR: aesCTR,
hmac: hmac,
ivSize: ivSize,
tagSize: tagSize,
prefix: key.OutputPrefix(),
variant: key.parameters.Variant(),
}, nil
}

func aadSizeInBits(associatedData []byte) []byte {
n := uint64(len(associatedData)) * 8
buf := make([]byte, 8)
binary.BigEndian.PutUint64(buf, n)
return buf
}

// Encrypt encrypts plaintext with associatedData.
//
// The plaintext is encrypted with an AES-CTR, then HMAC is computed over
// (associatedData || ciphertext || n) where n is associatedData's length
// in bits represented as a 64-bit big endian unsigned integer. The final
// ciphertext format is (IND-CPA ciphertext || mac).
func (a *fullAEAD) Encrypt(plaintext, associatedData []byte) ([]byte, error) {
ciphertext, err := a.aead.Encrypt(plaintext, associatedData)
ctSize := len(a.prefix) + a.ivSize + len(plaintext)
ciphertext := make([]byte, ctSize, ctSize+a.tagSize)
if n := copy(ciphertext, a.prefix); n != len(a.prefix) {
return nil, fmt.Errorf("aesctrhmac: failed to copy prefix")
}
ctNoPrefix, err := a.aesCTR.Encrypt(ciphertext[len(a.prefix):], plaintext)
if err != nil {
return nil, err
}
tag, err := a.hmac.ComputeMAC(associatedData, ctNoPrefix, aadSizeInBits(associatedData))
if err != nil {
return nil, err
}
return slices.Concat(a.prefix, ciphertext), nil
if len(tag) != a.tagSize {
return nil, errors.New("aesctrhmac: invalid tag size")
}
return append(ciphertext, tag...), nil
}

// Decrypt decrypts ciphertext with associatedData.
func (a *fullAEAD) Decrypt(ciphertext, associatedData []byte) ([]byte, error) {
if len(ciphertext) < len(a.prefix) {
return nil, fmt.Errorf("ciphertext with size %d is too short", len(ciphertext))
prefixSize := len(a.prefix)
if len(ciphertext) < prefixSize+a.ivSize+a.tagSize {
return nil, fmt.Errorf("aesctrhmac: ciphertext with size %d is too short", len(ciphertext))
}
prefix := ciphertext[:len(a.prefix)]
ciphertextNoPrefix := ciphertext[len(a.prefix):]
prefix := ciphertext[:prefixSize]
if !bytes.Equal(prefix, a.prefix) {
return nil, fmt.Errorf("ciphertext prefix does not match: got %x, want %x", prefix, a.prefix)
return nil, fmt.Errorf("aesctrhmac: ciphertext prefix does not match: got %x, want %x", prefix, a.prefix)
}
payload := ciphertext[prefixSize : len(ciphertext)-a.tagSize]
tag := ciphertext[len(ciphertext)-a.tagSize:]
if err := a.hmac.VerifyMAC(tag, associatedData, payload, aadSizeInBits(associatedData)); err != nil {
return nil, fmt.Errorf("aesctrhmac: %v", err)
}
return a.aead.Decrypt(ciphertextNoPrefix, associatedData)
return a.aesCTR.Decrypt(nil, payload)
}

// primitiveConstructor creates a [tink.AEAD] from a [key.Key].
Expand Down
36 changes: 22 additions & 14 deletions aead/aesctrhmac/key_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,9 @@ import (

"google.golang.org/protobuf/proto"
"github.com/tink-crypto/tink-go/v2/aead/subtle"
"github.com/tink-crypto/tink-go/v2/insecuresecretdataaccess"
"github.com/tink-crypto/tink-go/v2/keyset"
subtleMac "github.com/tink-crypto/tink-go/v2/mac/subtle"
"github.com/tink-crypto/tink-go/v2/secretdata"
"github.com/tink-crypto/tink-go/v2/subtle/random"
ctrpb "github.com/tink-crypto/tink-go/v2/proto/aes_ctr_go_proto"
aeadpb "github.com/tink-crypto/tink-go/v2/proto/aes_ctr_hmac_aead_go_proto"
Expand All @@ -44,28 +45,35 @@ type keyManager struct{}
// Primitive creates a [subtle.NewEncryptThenAuthenticate] primitive for the given
// serialized [aeadpb.AesCtrHmacAeadKey].
func (km *keyManager) Primitive(serializedKey []byte) (any, error) {
key := new(aeadpb.AesCtrHmacAeadKey)
if err := proto.Unmarshal(serializedKey, key); err != nil {
protoKey := new(aeadpb.AesCtrHmacAeadKey)
if err := proto.Unmarshal(serializedKey, protoKey); err != nil {
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: invalid key")
}
if err := km.validateKey(key); err != nil {
if err := km.validateKey(protoKey); err != nil {
return nil, err
}

ctr, err := subtle.NewAESCTR(key.GetAesCtrKey().GetKeyValue(), int(key.GetAesCtrKey().GetParams().GetIvSize()))
params, err := NewParameters(ParametersOpts{
AESKeySizeInBytes: int(len(protoKey.GetAesCtrKey().GetKeyValue())),
HMACKeySizeInBytes: int(len(protoKey.GetHmacKey().GetKeyValue())),
IVSizeInBytes: int(protoKey.GetAesCtrKey().GetParams().GetIvSize()),
TagSizeInBytes: int(protoKey.GetHmacKey().GetParams().GetTagSize()),
HashType: HashType(protoKey.GetHmacKey().GetParams().GetHash()),
Variant: VariantNoPrefix,
})
if err != nil {
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: cannot create new primitive: %v", err)
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: %s", err)
}

hmacKey := key.GetHmacKey()
hmac, err := subtleMac.NewHMAC(hmacKey.GetParams().GetHash().String(), hmacKey.GetKeyValue(), hmacKey.GetParams().GetTagSize())
key, err := NewKey(KeyOpts{
AESKeyBytes: secretdata.NewBytesFromData(protoKey.GetAesCtrKey().GetKeyValue(), insecuresecretdataaccess.Token{}),
HMACKeyBytes: secretdata.NewBytesFromData(protoKey.GetHmacKey().GetKeyValue(), insecuresecretdataaccess.Token{}),
Parameters: params,
})
if err != nil {
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: cannot create mac primitive, error: %v", err)
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: %s", err)
}

aead, err := subtle.NewEncryptThenAuthenticate(ctr, hmac, int(hmacKey.GetParams().GetTagSize()))
aead, err := newAEAD(key)
if err != nil {
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: cannot create encrypt then authenticate primitive, error: %v", err)
return nil, fmt.Errorf("aes_ctr_hmac_aead_key_manager: %s", err)
}
return aead, nil
}
Expand Down
74 changes: 8 additions & 66 deletions aead/subtle/aes_ctr.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,9 @@
package subtle

import (
"crypto/aes"
"crypto/cipher"
"fmt"
"github.com/tink-crypto/tink-go/v2/internal/aead"

// Placeholder for internal crypto/cipher allowlist, please ignore.
"github.com/tink-crypto/tink-go/v2/subtle/random"
)

const (
Expand All @@ -30,7 +27,7 @@ const (

// AESCTR is an implementation of AEAD interface.
type AESCTR struct {
key []byte
ctr *aead.AESCTR
IVSize int
}

Expand All @@ -39,76 +36,21 @@ type AESCTR struct {
// AES-128 or AES-256.
// ivSize specifies the size of the IV in bytes.
func NewAESCTR(key []byte, ivSize int) (*AESCTR, error) {
keySize := uint32(len(key))
if err := ValidateAESKeySize(keySize); err != nil {
return nil, fmt.Errorf("aes_ctr: %s", err)
}
if ivSize < AESCTRMinIVSize || ivSize > aes.BlockSize {
return nil, fmt.Errorf("aes_ctr: invalid IV size: %d", ivSize)
ctr, err := aead.NewAESCTR(key, ivSize)
if err != nil {
return nil, err
}
return &AESCTR{key: key, IVSize: ivSize}, nil
return &AESCTR{ctr: ctr, IVSize: ivSize}, nil
}

// Encrypt encrypts plaintext using AES in CTR mode.
// The resulting ciphertext consists of two parts:
// (1) the IV used for encryption and (2) the actual ciphertext.
func (a *AESCTR) Encrypt(plaintext []byte) ([]byte, error) {
if len(plaintext) > maxInt-a.IVSize {
return nil, fmt.Errorf("aes_ctr: plaintext too long")
}
iv := a.newIV()
stream, err := newCipher(a.key, iv)
if err != nil {
return nil, err
}

ciphertext := make([]byte, a.IVSize+len(plaintext))
if n := copy(ciphertext, iv); n != a.IVSize {
return nil, fmt.Errorf("aes_ctr: failed to copy IV (copied %d/%d bytes)", n, a.IVSize)
}

stream.XORKeyStream(ciphertext[a.IVSize:], plaintext)
return ciphertext, nil
return a.ctr.Encrypt(nil, plaintext)
}

// Decrypt decrypts ciphertext.
func (a *AESCTR) Decrypt(ciphertext []byte) ([]byte, error) {
if len(ciphertext) < a.IVSize {
return nil, fmt.Errorf("aes_ctr: ciphertext too short")
}

iv := ciphertext[:a.IVSize]
stream, err := newCipher(a.key, iv)
if err != nil {
return nil, err
}

plaintext := make([]byte, len(ciphertext)-a.IVSize)
stream.XORKeyStream(plaintext, ciphertext[a.IVSize:])
return plaintext, nil
}

// newIV creates a new IV for encryption.
func (a *AESCTR) newIV() []byte {
return random.GetRandomBytes(uint32(a.IVSize))
}

// newCipher creates a new AES-CTR cipher using the given key, IV and the crypto library.
func newCipher(key, iv []byte) (cipher.Stream, error) {
block, err := aes.NewCipher(key)
if err != nil {
return nil, fmt.Errorf("aes_ctr: failed to create block cipher, error: %v", err)
}

// If the IV is less than BlockSize bytes we need to pad it with zeros
// otherwise NewCTR will panic.
if len(iv) < aes.BlockSize {
paddedIV := make([]byte, aes.BlockSize)
if n := copy(paddedIV, iv); n != len(iv) {
return nil, fmt.Errorf("aes_ctr: failed to pad IV")
}
return cipher.NewCTR(block, paddedIV), nil
}

return cipher.NewCTR(block, iv), nil
return a.ctr.Decrypt(nil, ciphertext)
}
108 changes: 108 additions & 0 deletions internal/aead/aesctr.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aead

import (
"crypto/aes"
"crypto/cipher"
"fmt"

"github.com/tink-crypto/tink-go/v2/subtle/random"
)

const (
aesCTRMinIVSize = 12
)

// AESCTR is an implementation of IndCpa Interface.
type AESCTR struct {
key []byte
ivSize int
}

// NewAESCTR returns an instance of [AESCTR] unauthenticated encryption.
func NewAESCTR(key []byte, ivSize int) (*AESCTR, error) {
keySize := uint32(len(key))
if err := ValidateAESKeySize(keySize); err != nil {
return nil, fmt.Errorf("aes_ctr: %s", err)
}
if ivSize < aesCTRMinIVSize || ivSize > aes.BlockSize {
return nil, fmt.Errorf("aes_ctr: invalid IV size: %d", ivSize)
}
return &AESCTR{key: key, ivSize: ivSize}, nil
}

// Encrypt encrypts plaintext using AES in CTR mode.
func (a *AESCTR) Encrypt(dst, plaintext []byte) ([]byte, error) {
if len(plaintext) > maxInt-a.ivSize {
return nil, fmt.Errorf("aes_ctr: plaintext too long")
}
ctSize := len(plaintext) + a.ivSize
if len(dst) == 0 {
dst = make([]byte, ctSize)
}
if len(dst) < ctSize {
return nil, fmt.Errorf("aes_ctr: destination buffer too small (%d vs %d)", len(dst), ctSize)
}

iv := random.GetRandomBytes(uint32(a.ivSize))
stream, err := newCipher(a.key, iv)
if err != nil {
return nil, err
}
if n := copy(dst, iv); n != a.ivSize {
return nil, fmt.Errorf("aes_ctr: failed to copy IV (copied %d/%d bytes)", n, a.ivSize)
}
stream.XORKeyStream(dst[a.ivSize:], plaintext)
return dst, nil
}

// Decrypt decrypts ciphertext in the format (prefix || iv || ciphertext).
func (a *AESCTR) Decrypt(dst, ciphertext []byte) ([]byte, error) {
if len(ciphertext) < a.ivSize {
return nil, fmt.Errorf("aes_ctr: ciphertext too short")
}
ptSize := len(ciphertext) - a.ivSize
if len(dst) == 0 {
dst = make([]byte, ptSize)
}
if len(dst) < ptSize {
return nil, fmt.Errorf("aes_ctr: destination buffer too small (%d vs %d)", len(dst), ptSize)
}
stream, err := newCipher(a.key, ciphertext[:a.ivSize])
if err != nil {
return nil, err
}
stream.XORKeyStream(dst, ciphertext[a.ivSize:])
return dst, nil
}

func newCipher(key, iv []byte) (cipher.Stream, error) {
block, err := aes.NewCipher(key)
if err != nil {
return nil, fmt.Errorf("aes_ctr: failed to create block cipher, error: %v", err)
}

// If the IV is less than BlockSize bytes we need to pad it with zeros
// otherwise NewCTR will panic.
if len(iv) < aes.BlockSize {
paddedIV := make([]byte, aes.BlockSize)
if n := copy(paddedIV, iv); n != len(iv) {
return nil, fmt.Errorf("aes_ctr: failed to pad IV")
}
iv = paddedIV
}
return cipher.NewCTR(block, iv), nil
}
Loading

0 comments on commit e18b6d1

Please sign in to comment.