pqsafe/pkg/aead/aead.go

// Copyright (C) 2025-2026, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

// Package aead implements PQSAFE chunk AEAD: AES-256-GCM (FIPS profile)
// or ChaCha20-Poly1305 (software profile) over the chunked file payload.
//
// Nonce construction: 12 bytes = 8-byte nonce_seed (from manifest) ||
// 4-byte big-endian chunk index. The 8-byte seed is bound into the
// signed manifest so a sender cannot reuse the same nonce_seed across
// two different files without breaking the signature. Within one file,
// chunks have unique indices so nonces are unique.
//
// The DEK is rotated every 2^32 chunks (256 TiB) via HKDF re-key — but
// chunk size is 64 KiB and the largest realistic file is well under
// 256 TiB, so in practice one DEK protects one file.
package aead

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
	"fmt"

	"golang.org/x/crypto/chacha20poly1305"
)

// Profile chooses the bulk AEAD construction.
type Profile uint8

const (
	ProfileAESGCM           Profile = 1 // AES-256-GCM (FIPS)
	ProfileChaCha20Poly1305 Profile = 2 // ChaCha20-Poly1305 (software)
)

// String prints the canonical wire name.
func (p Profile) String() string {
	switch p {
	case ProfileAESGCM:
		return "AES-256-GCM"
	case ProfileChaCha20Poly1305:
		return "ChaCha20-Poly1305"
	default:
		return fmt.Sprintf("aead-profile(%d)", uint8(p))
	}
}

const (
	NonceSize    = 12
	NonceSeedLen = 8
	KeySize      = 32
	TagSize      = 16
)

// Errors.
var (
	ErrInvalidKey   = errors.New("aead: key must be 32 bytes")
	ErrInvalidNonce = errors.New("aead: nonce seed must be 8 bytes")
	ErrBadProfile   = errors.New("aead: unknown profile")
	ErrOpen         = errors.New("aead: authentication failed")
)

// Seal encrypts plaintext at the given chunk index. The full nonce is
// nonce_seed (8 bytes) || index (4 bytes, big-endian). The AAD binds
// the manifest hash + chunk index + profile so a chunk cannot be
// replayed into a different file or shuffled within the same file.
func Seal(profile Profile, key []byte, nonceSeed []byte, chunkIdx uint32, manifestHash []byte, plaintext []byte) ([]byte, error) {
	if len(key) != KeySize {
		return nil, ErrInvalidKey
	}
	if len(nonceSeed) != NonceSeedLen {
		return nil, ErrInvalidNonce
	}

	gcm, err := newAEAD(profile, key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, NonceSize)
	copy(nonce[:NonceSeedLen], nonceSeed)
	binary.BigEndian.PutUint32(nonce[NonceSeedLen:], chunkIdx)

	aad := makeAAD(profile, chunkIdx, manifestHash)
	return gcm.Seal(nil, nonce, plaintext, aad), nil
}

// Open inverts Seal. Returns ErrOpen if the authentication tag is bad
// or the AAD doesn't match (e.g., chunk replayed into a different
// manifest).
func Open(profile Profile, key []byte, nonceSeed []byte, chunkIdx uint32, manifestHash []byte, ciphertext []byte) ([]byte, error) {
	if len(key) != KeySize {
		return nil, ErrInvalidKey
	}
	if len(nonceSeed) != NonceSeedLen {
		return nil, ErrInvalidNonce
	}

	gcm, err := newAEAD(profile, key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, NonceSize)
	copy(nonce[:NonceSeedLen], nonceSeed)
	binary.BigEndian.PutUint32(nonce[NonceSeedLen:], chunkIdx)

	aad := makeAAD(profile, chunkIdx, manifestHash)
	pt, err := gcm.Open(nil, nonce, ciphertext, aad)
	if err != nil {
		return nil, ErrOpen
	}
	return pt, nil
}

// makeAAD binds (profile || chunk_idx || manifest_hash) into the AEAD's
// additional-authenticated-data. A chunk cipher-text replayed into a
// different manifest produces a tag mismatch.
func makeAAD(profile Profile, chunkIdx uint32, manifestHash []byte) []byte {
	aad := make([]byte, 1+4+len(manifestHash))
	aad[0] = byte(profile)
	binary.BigEndian.PutUint32(aad[1:5], chunkIdx)
	copy(aad[5:], manifestHash)
	return aad
}

func newAEAD(profile Profile, key []byte) (cipher.AEAD, error) {
	switch profile {
	case ProfileAESGCM:
		block, err := aes.NewCipher(key)
		if err != nil {
			return nil, err
		}
		return cipher.NewGCM(block)
	case ProfileChaCha20Poly1305:
		return chacha20poly1305.New(key)
	default:
		return nil, ErrBadProfile
	}
}