aead/xchacha20.go

// Copyright © 2023 Ory Corp
// SPDX-License-Identifier: Apache-2.0

package aead

import (
	"context"
	"crypto/cipher"
	cryptorand "crypto/rand"
	"encoding/base64"
	"fmt"
	"math"

	"golang.org/x/crypto/chacha20poly1305"

	"github.com/ory/x/errorsx"
)

var _ Cipher = (*XChaCha20Poly1305)(nil)

type (
	XChaCha20Poly1305 struct {
		d Dependencies
	}
)

func NewXChaCha20Poly1305(d Dependencies) *XChaCha20Poly1305 {
	return &XChaCha20Poly1305{d}
}

func (x *XChaCha20Poly1305) Encrypt(ctx context.Context, plaintext, additionalData []byte) (string, error) {
	key, err := encryptionKey(ctx, x.d, chacha20poly1305.KeySize)
	if err != nil {
		return "", err
	}

	aead, err := chacha20poly1305.NewX(key)
	if err != nil {
		return "", errorsx.WithStack(err)
	}

	// Make sure the size calculation does not overflow.
	if len(plaintext) > math.MaxInt-aead.NonceSize()-aead.Overhead() {
		return "", errorsx.WithStack(fmt.Errorf("plaintext too large"))
	}

	nonce := make([]byte, aead.NonceSize(), aead.NonceSize()+len(plaintext)+aead.Overhead())
	_, err = cryptorand.Read(nonce)
	if err != nil {
		return "", errorsx.WithStack(err)
	}

	ciphertext := aead.Seal(nonce, nonce, plaintext, additionalData)
	return base64.URLEncoding.EncodeToString(ciphertext), nil
}

func (x *XChaCha20Poly1305) Decrypt(ctx context.Context, ciphertext string, aad []byte) (plaintext []byte, err error) {
	msg, err := base64.URLEncoding.DecodeString(ciphertext)
	if err != nil {
		return nil, errorsx.WithStack(err)
	}

	if len(msg) < chacha20poly1305.NonceSizeX {
		return nil, errorsx.WithStack(fmt.Errorf("malformed ciphertext: too short"))
	}
	nonce, ciphered := msg[:chacha20poly1305.NonceSizeX], msg[chacha20poly1305.NonceSizeX:]

	keys, err := allKeys(ctx, x.d)
	if err != nil {
		return nil, errorsx.WithStack(err)
	}

	var aead cipher.AEAD
	for _, key := range keys {
		aead, err = chacha20poly1305.NewX(key)
		if err != nil {
			continue
		}
		plaintext, err = aead.Open(nil, nonce, ciphered, aad)
		if err == nil {
			return plaintext, nil
		}
	}

	return nil, errorsx.WithStack(err)
}