pqsafe/pkg/capsule/capsule.go

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

// Package capsule implements the PQSAFE recipient capsule: a per-
// recipient wrap of the per-file DEK under the recipient's X-Wing
// public key. X-Wing primitives are consumed from luxfi/zwing — the
// single canonical Lux implementation of the IETF
// draft-connolly-cfrg-xwing-kem construction.
//
// Wire format:
//
//	capsule := XWingCiphertext (zwing.XWingCiphertextSize bytes) ||
//	          AEAD-nonce      (12 bytes random per capsule)      ||
//	          AEAD(DEK; key=HKDF(ss, "PQSAFE-DEK-WRAP-v1"); aad=manifest_hash)
//
// The capsule binds to the manifest hash through AEAD AAD: a capsule
// from one manifest cannot be replayed into another manifest.
package capsule

import (
	"crypto/rand"
	"crypto/sha512"
	"errors"
	"fmt"

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

	"github.com/luxfi/zwing"

	"github.com/luxfi/pqsafe/pkg/ctx"
)

const (
	// WrapNonceSize is the AEAD nonce for the DEK wrap.
	WrapNonceSize = chacha20poly1305.NonceSize // 12
	// WrapAEADKeySize is the HKDF-derived wrap-key size.
	WrapAEADKeySize = 32
	// DEKSize is the file DEK length.
	DEKSize = 32
)

// Errors.
var (
	ErrBadDEKSize = errors.New("capsule: DEK must be 32 bytes")
	ErrUnwrap     = errors.New("capsule: unwrap failed (wrong recipient or tampered)")
	ErrBadCapsule = errors.New("capsule: malformed capsule bytes")
)

// Capsule is the on-wire form of a single recipient's wrapped DEK.
type Capsule struct {
	KEMCT []byte // zwing.XWingCiphertextSize bytes
	Nonce [WrapNonceSize]byte
	Wrap  []byte // AEAD(DEK; aad=manifest_hash) — 32 + 16 = 48 bytes
}

// MarshalBinary serializes a capsule.
func (c *Capsule) MarshalBinary() []byte {
	out := make([]byte, 0, len(c.KEMCT)+WrapNonceSize+len(c.Wrap))
	out = append(out, c.KEMCT...)
	out = append(out, c.Nonce[:]...)
	out = append(out, c.Wrap...)
	return out
}

// UnmarshalBinary parses a capsule from bytes.
func UnmarshalBinary(b []byte) (*Capsule, error) {
	if len(b) < zwing.XWingCiphertextSize+WrapNonceSize+DEKSize+16 {
		return nil, ErrBadCapsule
	}
	c := &Capsule{KEMCT: make([]byte, zwing.XWingCiphertextSize)}
	copy(c.KEMCT, b[:zwing.XWingCiphertextSize])
	copy(c.Nonce[:], b[zwing.XWingCiphertextSize:zwing.XWingCiphertextSize+WrapNonceSize])
	c.Wrap = append([]byte(nil), b[zwing.XWingCiphertextSize+WrapNonceSize:]...)
	return c, nil
}

// Seal wraps a file DEK to one recipient.
func Seal(recipient *zwing.XWingPublicKey, dek, manifestHash []byte) (*Capsule, error) {
	if len(dek) != DEKSize {
		return nil, ErrBadDEKSize
	}

	kemCT, ss, err := zwing.XWingEncapsulate(rand.Reader, recipient)
	if err != nil {
		return nil, err
	}

	wrapKey := make([]byte, WrapAEADKeySize)
	kdf := hkdf.New(sha512.New384, ss[:], nil, []byte(ctx.DEKWrap))
	if _, err := kdf.Read(wrapKey); err != nil {
		return nil, err
	}

	aead, err := chacha20poly1305.New(wrapKey)
	if err != nil {
		return nil, err
	}

	c := &Capsule{KEMCT: kemCT}
	if _, err := rand.Read(c.Nonce[:]); err != nil {
		return nil, err
	}
	c.Wrap = aead.Seal(nil, c.Nonce[:], dek, manifestHash)
	return c, nil
}

// Open unwraps the DEK using the recipient's secret key. The
// manifest_hash MUST match the bytes used in Seal.
func (c *Capsule) Open(recipientSec *zwing.XWingPrivateKey, manifestHash []byte) ([]byte, error) {
	ss, err := zwing.XWingDecapsulate(recipientSec, c.KEMCT)
	if err != nil {
		return nil, fmt.Errorf("capsule: zwing.XWingDecapsulate: %w", err)
	}

	wrapKey := make([]byte, WrapAEADKeySize)
	kdf := hkdf.New(sha512.New384, ss[:], nil, []byte(ctx.DEKWrap))
	if _, err := kdf.Read(wrapKey); err != nil {
		return nil, err
	}

	aead, err := chacha20poly1305.New(wrapKey)
	if err != nil {
		return nil, err
	}
	dek, err := aead.Open(nil, c.Nonce[:], c.Wrap, manifestHash)
	if err != nil {
		return nil, ErrUnwrap
	}
	if len(dek) != DEKSize {
		return nil, ErrUnwrap
	}
	return dek, nil
}