DateKeys/agewrap/agewrap.go

// Package agewrap holds the age recipients and identities that DateKeys wraps
// around standard age files (spec §28-§37), plus the structural stanza rules
// every DateKeys age file must satisfy.
//
// The cryptography is age, tlock and drand's BLS verification. This package
// adds only the rules of the protocol:
//
//   - OUTER_TIME_AGE holds exactly one tlock stanza for the expected round and
//     the pinned chain hash (spec §32, §35, §63 step 11).
//   - PAYLOAD_AGE holds exactly one X25519 stanza, for R_PAYLOAD (spec §29,
//     §63 step 17).
//   - INNER_ACCESS_AGE holds one or more stanzas, all X25519, one per
//     recipient (spec §33, §63 step 13).
//
// The rules are enforced inside Identity.Unwrap, which age calls with the
// complete set of stanzas of the file, so that no file is accepted just
// because age managed to unwrap a file key (spec §27, §63). The same checks
// are exposed for the pre-unlock inspection, which reads the stanzas through a
// probe identity without decrypting anything or touching secrets.
package agewrap

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	"filippo.io/age"
	"github.com/drand/drand/v2/common"
	"github.com/drand/drand/v2/crypto"
	"github.com/drand/kyber"
	"github.com/drand/tlock"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/codec/bech32"
	"g.activething.com/go/DateKeys/profile"
	"g.activething.com/go/DateKeys/provider"
)

// Stanza types of V1.
const (
	StanzaTLock  = "tlock"
	StanzaX25519 = "X25519"
)

// FileKeySize is the size of every age file key: FK_PAYLOAD, FK_ACCESS and
// FK_TIME (spec §28).
const FileKeySize = 16

// ---------------------------------------------------------------------------
// Structural rules

// CheckTimeStanzas enforces the OUTER_TIME_AGE rule: exactly one stanza, of
// type tlock, whose round is the DateKey round and whose chain hash is the one
// of the pinned profile (spec §32, §35, §63 steps 5, 8 and 11).
func CheckTimeStanzas(stanzas []*age.Stanza, p *profile.Profile, round uint64) error {
	if len(stanzas) != 1 {
		return fmt.Errorf("agewrap: OUTER_TIME_AGE has %d stanzas, want exactly one tlock stanza: %w", len(stanzas), datekeys.ErrPolicyStructureMismatch)
	}
	s := stanzas[0]
	if s.Type != StanzaTLock {
		return fmt.Errorf("agewrap: OUTER_TIME_AGE stanza type %q, want %q: %w", s.Type, StanzaTLock, datekeys.ErrPolicyStructureMismatch)
	}
	if len(s.Args) != 2 {
		return fmt.Errorf("agewrap: tlock stanza has %d arguments, want 2: %w", len(s.Args), datekeys.ErrPolicyStructureMismatch)
	}
	if want := strconv.FormatUint(round, 10); s.Args[0] != want {
		return fmt.Errorf("agewrap: tlock stanza round %q, DateKey round %s: %w", s.Args[0], want, datekeys.ErrRoundMismatch)
	}
	if want := p.ChainHashHex(); s.Args[1] != want {
		return fmt.Errorf("agewrap: tlock stanza chain hash %q, pinned profile %s uses %s: %w", s.Args[1], p.ID, want, datekeys.ErrProfileMismatch)
	}
	return nil
}

// CheckPayloadStanzas enforces the PAYLOAD_AGE rule: exactly one stanza, of
// type X25519 (spec §29, §63 steps 6 and 17).
func CheckPayloadStanzas(stanzas []*age.Stanza) error {
	if len(stanzas) != 1 {
		return fmt.Errorf("agewrap: PAYLOAD_AGE has %d stanzas, want exactly one X25519 stanza: %w", len(stanzas), datekeys.ErrPolicyStructureMismatch)
	}
	if t := stanzas[0].Type; t != StanzaX25519 {
		return fmt.Errorf("agewrap: PAYLOAD_AGE stanza type %q, want %q: %w", t, StanzaX25519, datekeys.ErrPolicyStructureMismatch)
	}
	return nil
}

// CheckAccessStanzas enforces the INNER_ACCESS_AGE rule: one or more stanzas,
// all of type X25519 (spec §33, §36, §63 step 13). Two stanzas with the same
// ephemeral share would be two stanzas for one recipient and are rejected.
func CheckAccessStanzas(stanzas []*age.Stanza) error {
	if len(stanzas) == 0 {
		return fmt.Errorf("agewrap: INNER_ACCESS_AGE has no stanzas: %w", datekeys.ErrPolicyStructureMismatch)
	}
	seen := make(map[string]bool, len(stanzas))
	for i, s := range stanzas {
		if s.Type != StanzaX25519 {
			return fmt.Errorf("agewrap: INNER_ACCESS_AGE stanza %d has type %q, want %q: %w", i, s.Type, StanzaX25519, datekeys.ErrPolicyStructureMismatch)
		}
		if len(s.Args) == 1 {
			if seen[s.Args[0]] {
				return fmt.Errorf("agewrap: INNER_ACCESS_AGE stanza %d repeats an ephemeral share: %w", i, datekeys.ErrPolicyStructureMismatch)
			}
			seen[s.Args[0]] = true
		}
	}
	return nil
}

// ---------------------------------------------------------------------------
// Inspection probe

var errProbe = errors.New("agewrap: probe finished")

// probe records the stanzas age hands to Unwrap and stops decryption with an
// error that does not wrap age.ErrIncorrectIdentity, so age returns it as is.
type probe struct{ stanzas []*age.Stanza }

func (p *probe) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	p.stanzas = cloneStanzas(stanzas)
	return nil, errProbe
}

// Stanzas parses the age header at the start of r with age itself and returns
// its recipient stanzas. It decrypts nothing and uses no secret: the header is
// extracted with age.ExtractHeader and handed to age.DecryptHeader with a
// probe identity (spec §27, §63 steps 5 and 6).
//
// The result is structural. Its authenticity is only established when the
// header MAC is verified while opening the file (spec §27).
func Stanzas(r io.Reader) ([]*age.Stanza, error) {
	hdr, err := age.ExtractHeader(r)
	if err != nil {
		return nil, fmt.Errorf("agewrap: not a valid age file: %v: %w", err, datekeys.ErrIntegrity)
	}
	var p probe
	if _, err := age.DecryptHeader(hdr, &p); !errors.Is(err, errProbe) {
		return nil, fmt.Errorf("agewrap: unexpected result inspecting the age header: %v: %w", err, datekeys.ErrIntegrity)
	}
	return p.stanzas, nil
}

func cloneStanzas(in []*age.Stanza) []*age.Stanza {
	out := make([]*age.Stanza, len(in))
	for i, s := range in {
		out[i] = &age.Stanza{Type: s.Type, Args: append([]string(nil), s.Args...), Body: bytes.Clone(s.Body)}
	}
	return out
}

// ---------------------------------------------------------------------------
// tlock recipient and identity (OUTER_TIME_AGE)

// TimeRecipient wraps the file key with tlock for one round of a pinned
// profile and emits the stanza "tlock <round> <chainhash>", byte-compatible
// with the stanza of the tlock library and the tle CLI (spec §32, §35). It
// uses only the exported core of tlock: TimeLock and CiphertextToBytes.
type TimeRecipient struct {
	chainHash string
	round     uint64
	scheme    *crypto.Scheme
	key       kyber.Point
}

var _ age.RecipientWithLabels = (*TimeRecipient)(nil)

// NewTimeRecipient returns the tlock recipient of round under p, using only
// the pinned parameters of p (strict mode, spec §35).
func NewTimeRecipient(p *profile.Profile, round uint64) (*TimeRecipient, error) {
	scheme, key, err := pinned(p)
	if err != nil {
		return nil, err
	}
	if round == 0 || round > p.MaxRound() {
		return nil, fmt.Errorf("agewrap: round %d outside the range of %s: %w", round, p.ID, datekeys.ErrDateKeyInvalid)
	}
	return &TimeRecipient{chainHash: p.ChainHashHex(), round: round, scheme: scheme, key: key}, nil
}

// Wrap implements age.Recipient.
func (r *TimeRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
	ct, err := tlock.TimeLock(*r.scheme, r.key, r.round, fileKey)
	if err != nil {
		return nil, fmt.Errorf("agewrap: tlock: %w", err)
	}
	body, err := tlock.CiphertextToBytes(*r.scheme, ct)
	if err != nil {
		return nil, fmt.Errorf("agewrap: tlock ciphertext: %w", err)
	}
	return []*age.Stanza{{
		Type: StanzaTLock,
		Args: []string{strconv.FormatUint(r.round, 10), r.chainHash},
		Body: body,
	}}, nil
}

// WrapWithLabels implements age.RecipientWithLabels with a random label, so
// that age refuses to mix this recipient with any other one in the same file:
// OUTER_TIME_AGE must hold exactly one tlock stanza.
func (r *TimeRecipient) WrapWithLabels(fileKey []byte) ([]*age.Stanza, []string, error) {
	s, err := r.Wrap(fileKey)
	if err != nil {
		return nil, nil, err
	}
	var label [16]byte
	_, _ = rand.Read(label[:]) // never fails since Go 1.24
	return s, []string{"datekeys-tlock-" + hex.EncodeToString(label[:])}, nil
}

// TimeIdentity opens OUTER_TIME_AGE under the strict rules of spec §35 and
// §63 step 11. Unwrap validates the complete stanza set, verifies the release
// locally and calls tlock.TimeUnlock, which verifies the beacon again before
// decrypting. Every failure keeps its own normative error: none is turned
// into "too early".
type TimeIdentity struct {
	profile *profile.Profile
	round   uint64
	release provider.Release
	scheme  *crypto.Scheme
	key     kyber.Point
}

var _ age.Identity = (*TimeIdentity)(nil)

// NewTimeIdentity returns the identity that opens OUTER_TIME_AGE for round
// with release.
func NewTimeIdentity(p *profile.Profile, round uint64, release provider.Release) (*TimeIdentity, error) {
	scheme, key, err := pinned(p)
	if err != nil {
		return nil, err
	}
	return &TimeIdentity{profile: p.Clone(), round: round, release: release, scheme: scheme, key: key}, nil
}

// Unwrap implements age.Identity.
func (i *TimeIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckTimeStanzas(stanzas, i.profile, i.round); err != nil {
		return nil, err
	}
	if err := provider.Verify(i.profile, provider.Condition{Round: i.round}, i.release); err != nil {
		return nil, err
	}
	ct, err := tlock.BytesToCiphertext(*i.scheme, stanzas[0].Body)
	if err != nil {
		return nil, fmt.Errorf("agewrap: malformed tlock stanza body: %v: %w", err, datekeys.ErrIntegrity)
	}
	beacon := common.Beacon{Round: i.release.Round, Signature: i.release.Signature}
	fileKey, err := tlock.TimeUnlock(*i.scheme, i.key, beacon, ct)
	if err != nil {
		return nil, fmt.Errorf("agewrap: tlock unwrap failed: %v: %w", err, datekeys.ErrIntegrity)
	}
	if len(fileKey) != FileKeySize {
		return nil, fmt.Errorf("agewrap: tlock stanza wraps a %d-byte file key: %w", len(fileKey), datekeys.ErrIntegrity)
	}
	return fileKey, nil
}

func pinned(p *profile.Profile) (*crypto.Scheme, kyber.Point, error) {
	scheme, err := p.DrandScheme()
	if err != nil {
		return nil, nil, err
	}
	key := scheme.KeyGroup.Point()
	if err := key.UnmarshalBinary(p.PublicKey); err != nil {
		return nil, nil, fmt.Errorf("agewrap: pinned public key of %s: %v: %w", p.ID, err, datekeys.ErrUnknownProfile)
	}
	if key.Equal(key.Null()) {
		return nil, nil, fmt.Errorf("agewrap: pinned public key of %s is the identity element: %w", p.ID, datekeys.ErrUnknownProfile)
	}
	return scheme, key, nil
}

// ---------------------------------------------------------------------------
// X25519 identities (PAYLOAD_AGE and INNER_ACCESS_AGE)

// PayloadIdentity opens PAYLOAD_AGE with I_PAYLOAD (spec §29, §30.1, §63 step
// 17). It rejects the file unless it holds exactly one X25519 stanza and that
// stanza is for R_PAYLOAD.
type PayloadIdentity struct {
	id *age.X25519Identity
}

var _ age.Identity = (*PayloadIdentity)(nil)

// NewPayloadIdentity returns the identity for the raw 32-byte I_PAYLOAD.
func NewPayloadIdentity(raw []byte) (*PayloadIdentity, error) {
	id, err := X25519IdentityFromRaw(raw)
	if err != nil {
		return nil, err
	}
	return &PayloadIdentity{id: id}, nil
}

// Unwrap implements age.Identity.
func (i *PayloadIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckPayloadStanzas(stanzas); err != nil {
		return nil, err
	}
	fileKey, err := i.id.Unwrap(stanzas)
	if errors.Is(err, age.ErrIncorrectIdentity) {
		// CONTROL_A + PAYLOAD_AGE_B: I_PAYLOAD_A cannot unwrap FK_PAYLOAD_B (spec §30.1).
		return nil, fmt.Errorf("agewrap: PAYLOAD_AGE is not encrypted to this control's R_PAYLOAD: %w", datekeys.ErrIntegrity)
	}
	if err != nil {
		return nil, fmt.Errorf("agewrap: malformed PAYLOAD_AGE stanza: %v: %w", err, datekeys.ErrIntegrity)
	}
	return fileKey, nil
}

// AccessIdentity opens INNER_ACCESS_AGE with the caller's X25519 identities,
// including the one of a portable .dkk (spec §33, §38, §63 step 13). It
// validates the complete stanza set first, and rejects the file if any
// identity unwraps more than one stanza, which would be two stanzas for the
// same recipient.
type AccessIdentity struct {
	ids []age.Identity
}

var _ age.Identity = (*AccessIdentity)(nil)

// NewAccessIdentity returns an AccessIdentity trying every identity of ids.
// Nil entries are dropped; at least one identity is required.
func NewAccessIdentity(ids ...age.Identity) (*AccessIdentity, error) {
	var clean []age.Identity
	for _, id := range ids {
		if id != nil {
			clean = append(clean, id)
		}
	}
	if len(clean) == 0 {
		return nil, fmt.Errorf("agewrap: time_and_key needs an access identity: %w", datekeys.ErrAccessRequired)
	}
	return &AccessIdentity{ids: clean}, nil
}

// Unwrap implements age.Identity. The codes follow spec §63 step 13, in this
// order: the stanza rules (ErrPolicyStructureMismatch); a malformed X25519
// stanza, found by the first identity already, because age checks the form
// of a stanza before any key agreement (ErrIntegrity); an identity that
// unwraps more than one stanza, even if another one unwraps exactly one
// (ErrPolicyStructureMismatch); no identity that unwraps any
// (ErrAccessInvalid). Every identity is tried against every stanza, so that
// the result does not depend on the order of the identities (spec §69.1).
func (a *AccessIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckAccessStanzas(stanzas); err != nil {
		return nil, err
	}
	var fileKey []byte
	for _, id := range a.ids {
		matches := 0
		for _, s := range stanzas {
			fk, err := id.Unwrap([]*age.Stanza{s})
			if errors.Is(err, age.ErrIncorrectIdentity) {
				continue
			}
			if err != nil {
				return nil, fmt.Errorf("agewrap: malformed INNER_ACCESS_AGE stanza: %v: %w", err, datekeys.ErrIntegrity)
			}
			matches++
			if fileKey == nil {
				fileKey = fk
			}
		}
		if matches > 1 {
			return nil, fmt.Errorf("agewrap: one identity opens %d INNER_ACCESS_AGE stanzas, want one per recipient: %w", matches, datekeys.ErrPolicyStructureMismatch)
		}
	}
	if fileKey == nil {
		return nil, fmt.Errorf("agewrap: no supplied identity is a recipient of INNER_ACCESS_AGE: %w", datekeys.ErrAccessInvalid)
	}
	return fileKey, nil
}

// ---------------------------------------------------------------------------
// Raw X25519 keys

// X25519IdentityFromRaw converts 32 raw identity bytes, the canonical form
// inside CONTROL_CBOR and .dkk (spec §31, §38), to an age identity.
func X25519IdentityFromRaw(raw []byte) (*age.X25519Identity, error) {
	if len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: X25519 identity is %d bytes, want 32: %w", len(raw), datekeys.ErrIntegrity)
	}
	s, err := bech32.Encode("AGE-SECRET-KEY-", raw)
	if err != nil {
		return nil, fmt.Errorf("agewrap: encode identity: %v: %w", err, datekeys.ErrIntegrity)
	}
	id, err := age.ParseX25519Identity(strings.ToUpper(s))
	if err != nil {
		return nil, fmt.Errorf("agewrap: parse identity: %v: %w", err, datekeys.ErrIntegrity)
	}
	return id, nil
}

// RawX25519Identity returns the 32 raw bytes of an age X25519 identity.
func RawX25519Identity(id *age.X25519Identity) ([]byte, error) {
	hrp, raw, err := bech32.Decode(id.String())
	if err != nil || hrp != "AGE-SECRET-KEY-" || len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: unexpected age identity encoding")
	}
	return raw, nil
}

// RawX25519Recipient returns the 32 raw bytes of an age X25519 recipient.
func RawX25519Recipient(r *age.X25519Recipient) ([]byte, error) {
	hrp, raw, err := bech32.Decode(r.String())
	if err != nil || hrp != "age" || len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: unexpected age recipient encoding")
	}
	return raw, nil
}