DateKeys/capsule/encrypt.go

package capsule

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"math/big"
	"time"

	"filippo.io/age"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/accesskey"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/profile"
)

// EncryptOptions configures Encrypt.
type EncryptOptions struct {
	// Profile is the pinned Provider Profile. Required.
	Profile *profile.Profile
	// UnlockAt is the requested instant. It resolves locally to the first
	// round at or after it (spec §15) and must be after Now.
	UnlockAt time.Time
	// Policy is time_only or time_and_key (spec §25).
	Policy Policy
	// Recipients are the X25519 recipients of known holders, for
	// time_and_key (spec §33, §37, §39). Only *age.X25519Recipient is
	// accepted: INNER_ACCESS_AGE must hold X25519 stanzas only. Each must be
	// canonical and not of low order, and none may be listed twice.
	Recipients []age.Recipient
	// NewPortableKey generates a fresh I_ACCESS for this capsule only and
	// returns it as a .dkk (spec §38). An I_ACCESS is never reused: Encrypt
	// accepts no existing one. The recipients and the portable key are the
	// credentials of the capsule: from 1 to 16 (spec §39).
	NewPortableKey bool
	// Length is L, the exact number of bytes src delivers, at most
	// MaxPayloadLength. It is sealed in the control before the payload is
	// written, so it must be known in advance: a source of unknown length can
	// be copied to a temporary file first (spec §29.1, §62.1 rule 6). If src
	// delivers another number of bytes, Encrypt fails.
	Length int64
	// Padding is the padding rule of the payload, Bloque256 or Reforzado.
	// Zero means Reforzado, the default of spec §29.1.
	Padding Padding
	// Critical and Noncritical are the PUBLIC_HEADER extensions (visible to
	// anyone holding the .dkc).
	Critical, Noncritical []extension.Extension
	// ControlCritical and ControlNoncritical are the CONTROL_CBOR extensions,
	// sealed with the control.
	ControlCritical, ControlNoncritical []extension.Extension
	// Now is the clock. Required: no package of this module reads the wall
	// clock on its own.
	Now func() time.Time
}

// Result describes a capsule written by Encrypt.
type Result struct {
	DateKey   datekey.DateKey
	UnlockAt  time.Time // effective round time, never before the requested instant
	CapsuleID [CapsuleIDSize]byte
	// Format is the format written, always Format2. Length is L, Padding the
	// padding rule and PaddedLength P = rule(L), the length of the plaintext
	// of PAYLOAD_AGE (spec §29.1).
	Format       Format
	Length       uint64
	Padding      Padding
	PaddedLength uint64
	// PortableKey is the .dkk generated when NewPortableKey is set. Encode it
	// with accesskey.Encode and treat it as a sensitive capability.
	PortableKey *accesskey.AccessKey
}

// Encrypt writes a format 2 .dkc for the content read from src (spec §61 for
// time_only, §62 for time_and_key, §62.1). It needs no network: the round is
// resolved locally and tlock uses only the pinned public key.
//
// PAYLOAD_AGE is streamed after the small, in-memory SEALED_CONTROL, so the
// content is never held in memory. Its plaintext is the content followed by
// zeros up to P = rule(L) (spec §29.1). On error dst may hold a partial
// capsule that must be discarded and never presented as a capsule (spec
// §62.1 rule 9).
//
// Before and after writing, Encrypt checks its own output with the rules of
// the reader (spec §62.1 rule 11): PUBLIC_HEADER and CONTROL_CBOR decode,
// INNER_ACCESS_AGE holds 16 X25519 stanzas with distinct shares and the
// portable key opens exactly one, the plaintext handed to age is P bytes and
// PAYLOAD_AGE has the length P gives, and I_PAYLOAD opens its header.
//
// The extensions of opts are written as given, once they pass the rules of
// spec §54. Encrypt takes no extension.Registry: the application writes a
// registered extension only in the objects and arrays it is registered for
// (spec §72).
func Encrypt(dst io.Writer, src io.Reader, opts EncryptOptions) (*Result, error) {
	p := opts.Profile
	if p == nil {
		return nil, errors.New("capsule: EncryptOptions.Profile is required")
	}
	if opts.Now == nil {
		return nil, errors.New("capsule: EncryptOptions.Now is required")
	}
	if err := p.Validate(); err != nil {
		return nil, err
	}
	if !opts.UnlockAt.After(opts.Now()) {
		return nil, fmt.Errorf("capsule: unlock time %s is not in the future", opts.UnlockAt.UTC().Format(time.RFC3339Nano))
	}
	if opts.Length < 0 {
		return nil, fmt.Errorf("capsule: EncryptOptions.Length %d is negative", opts.Length)
	}
	code := opts.Padding
	if code == 0 {
		code = Reforzado
	}
	length := uint64(opts.Length)
	padded, err := PaddedLength(length, code)
	if err != nil {
		return nil, err
	}

	// Step 2: resolve the DateKey locally.
	dk, err := datekey.Resolve(p, opts.UnlockAt)
	if err != nil {
		return nil, err
	}
	unlock := dk.UnlockAt(p)
	// Spec §17: round_time(round) >= requested_unlock_at, never earlier.
	if unlock.Before(opts.UnlockAt) {
		return nil, fmt.Errorf("capsule: resolved round %d opens before the requested time: %w", dk.Round, datekeys.ErrRoundMismatch)
	}

	credentials, portable, err := accessRecipients(opts)
	if err != nil {
		return nil, err
	}
	var portableRaw []byte
	if portable != nil {
		if portableRaw, err = agewrap.RawX25519Identity(portable); err != nil {
			return nil, err
		}
		defer clear(portableRaw)
	}

	// Step 3: capsule_id, 16 random bytes (spec §21).
	var capsuleID [CapsuleIDSize]byte
	_, _ = rand.Read(capsuleID[:]) // never fails since Go 1.24

	// Step 4: I_PAYLOAD, a fresh X25519 identity (spec §29).
	payloadID, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, err
	}
	payloadRaw, err := agewrap.RawX25519Identity(payloadID)
	if err != nil {
		return nil, err
	}
	defer clear(payloadRaw)

	// Step 6 of spec §62: the 16 recipients of INNER_ACCESS_AGE, the
	// credentials and a dummy in each slot left, in a random order.
	var access []age.Recipient
	if opts.Policy == TimeAndKey {
		if access, err = fillSlots(credentials); err != nil {
			return nil, err
		}
	}

	// Step 5 (7 of spec §62): PUBLIC_HEADER.
	header := &Header{CapsuleID: capsuleID, DateKey: dk, Policy: opts.Policy, Critical: opts.Critical, Noncritical: opts.Noncritical}
	headerBytes, err := EncodeHeader(header)
	if err != nil {
		return nil, err
	}
	if err := selfCheckHeader(headerBytes); err != nil {
		return nil, err
	}

	timeRecipient, err := agewrap.NewTimeRecipient(p, dk.Round)
	if err != nil {
		return nil, err
	}
	seal := func(control []byte) ([]byte, error) {
		plaintext := control
		if opts.Policy == TimeAndKey {
			// INNER_ACCESS_AGE: FK_ACCESS wrapped for the 16 recipients.
			innerAge, err := encryptAll(control, access...)
			if err != nil {
				return nil, err
			}
			if err := selfCheckInner(innerAge, control, portable); err != nil {
				return nil, err
			}
			plaintext = innerAge
		}
		// OUTER_TIME_AGE: FK_TIME wrapped with tlock for the DateKey round.
		return encryptAll(plaintext, timeRecipient)
	}

	// Steps 6 to 10 (8 to 13 of spec §62). PRELUDE carries
	// SEALED_CONTROL_LEN and header_binding covers PRELUDE, so the length is
	// measured first by sealing a control of identical size with a zero
	// binding and a zero identity: the length of a version 2 control does not
	// depend on them, on L or on the padding code (spec §62.1 rule 7). age
	// output lengths depend only on plaintext length and stanza shapes; the
	// real seal is checked to have the same length.
	ctrl := &Control{
		Critical: opts.ControlCritical, Noncritical: opts.ControlNoncritical,
		PayloadLength: length, Padding: code,
	}
	draft, err := EncodeControl(ctrl, Format2)
	if err != nil {
		return nil, err
	}
	draftSealed, err := seal(draft)
	if err != nil {
		return nil, err
	}
	if len(draftSealed) > MaxSealedControlLen {
		return nil, fmt.Errorf("capsule: SEALED_CONTROL of %d bytes exceeds %d: %w", len(draftSealed), MaxSealedControlLen, datekeys.ErrIntegrity)
	}
	prelude := Prelude{Format: Format2, PublicHeaderLen: uint32(len(headerBytes)), SealedControlLen: uint32(len(draftSealed))}
	preludeBytes := prelude.Bytes()

	// header_binding = SHA-256(PRELUDE || PUBLIC_HEADER_BYTES).
	ctrl.HeaderBinding = HeaderBinding(preludeBytes, headerBytes)
	copy(ctrl.PayloadIdentity[:], payloadRaw)
	defer clear(ctrl.PayloadIdentity[:])

	// CONTROL_CBOR, schema version 2, with L and the padding code.
	controlBytes, err := EncodeControl(ctrl, Format2)
	if err != nil {
		return nil, err
	}
	defer clear(controlBytes)
	if err := selfCheckControl(controlBytes); err != nil {
		return nil, err
	}

	// SEALED_CONTROL = OUTER_TIME_AGE.
	sealed, err := seal(controlBytes)
	if err != nil {
		return nil, err
	}
	if len(sealed) != len(draftSealed) {
		return nil, fmt.Errorf("capsule: internal error: SEALED_CONTROL is %d bytes, measured %d", len(sealed), len(draftSealed))
	}

	// PRELUDE || PUBLIC_HEADER || SEALED_CONTROL || PAYLOAD_AGE.
	digest := sha256.New()
	w := io.MultiWriter(dst, digest)
	for _, b := range [][]byte{preludeBytes[:], headerBytes, sealed} {
		if _, err := w.Write(b); err != nil {
			return nil, err
		}
	}
	// PAYLOAD_AGE, a standard age file for R_PAYLOAD (FK_PAYLOAD is
	// generated by age): the content and its padding, streamed.
	payload := &payloadWriter{w: w}
	aw, err := age.Encrypt(payload, payloadID.Recipient())
	if err != nil {
		return nil, err
	}
	if err := writeContent(aw, src, opts.Length, padded); err != nil {
		return nil, err
	}
	if err := aw.Close(); err != nil {
		return nil, err
	}
	if err := selfCheckPayload(payload, payloadRaw, padded); err != nil {
		return nil, err
	}

	res := &Result{DateKey: dk, UnlockAt: unlock, CapsuleID: capsuleID, Format: Format2, Length: length, Padding: code, PaddedLength: padded}
	if portable != nil {
		// The portable identity as 32 raw bytes in a .dkk (§62 step 17).
		k := &accesskey.AccessKey{
			CapsuleID:    capsuleID,
			Type:         accesskey.TypeX25519,
			Material:     bytes.Clone(portableRaw),
			Verification: &accesskey.Verification{CapsuleDigest: digest.Sum(nil)},
		}
		_, _ = rand.Read(k.CredentialID[:]) // spec §42; never fails since Go 1.24
		res.PortableKey = k
	}
	return res, nil
}

// writeContent writes to the age writer aw exactly length bytes of src, then
// the zeros of the padding up to padded bytes (spec §29.1). A source that
// delivers fewer or more than length bytes is an error: the capsule would
// fail at step 17, after the date, when it can no longer be repaired (spec
// §62.1 rule 6).
func writeContent(aw io.Writer, src io.Reader, length int64, padded uint64) error {
	n, err := io.CopyN(aw, src, length)
	if err == io.EOF {
		return fmt.Errorf("capsule: the source ended after %d bytes, and EncryptOptions.Length is %d", n, length)
	}
	if err != nil {
		return err
	}
	var more [1]byte
	switch _, err := io.ReadFull(src, more[:]); {
	case err == nil:
		return fmt.Errorf("capsule: the source delivers more than the %d bytes of EncryptOptions.Length", length)
	case err != io.EOF:
		return err
	}
	zeros := make([]byte, min(padded-uint64(length), 16<<10))
	for left := padded - uint64(length); left > 0; {
		k := min(left, uint64(len(zeros)))
		if _, err := aw.Write(zeros[:k]); err != nil {
			return err
		}
		left -= k
	}
	return nil
}

// payloadWriter counts the bytes of PAYLOAD_AGE and keeps the first ones,
// where its age header is, for the self-check.
type payloadWriter struct {
	w    io.Writer
	n    uint64
	head []byte
}

// payloadHeadSize bounds the bytes kept: an age header with one X25519
// stanza is 168 bytes.
const payloadHeadSize = 1 << 10

func (p *payloadWriter) Write(b []byte) (int, error) {
	if room := payloadHeadSize - len(p.head); room > 0 {
		p.head = append(p.head, b[:min(room, len(b))]...)
	}
	n, err := p.w.Write(b)
	p.n += uint64(n)
	return n, err
}

// selfCheckHeader decodes PUBLIC_HEADER with the reader's decoder before
// anything is sealed or written: a capsule whose header the reader rejects
// would be unusable (spec §62.1 rule 11, §72).
func selfCheckHeader(b []byte) error {
	if _, err := DecodeHeader(b); err != nil {
		return fmt.Errorf("capsule: self-check: the reader rejects this PUBLIC_HEADER: %w", err)
	}
	return nil
}

// selfCheckControl decodes CONTROL_CBOR with the reader's decoder before it is
// sealed. A control that the reader rejects would only be found at step 14
// of spec §63, after the unlock, when the capsule can no longer be repaired.
func selfCheckControl(b []byte) error {
	c, err := DecodeControl(b, Format2)
	if err != nil {
		return fmt.Errorf("capsule: self-check: the reader rejects this CONTROL_CBOR: %w", err)
	}
	clear(c.PayloadIdentity[:])
	return nil
}

// selfCheckInner checks INNER_ACCESS_AGE with the rules of the reader: 16
// X25519 stanzas with distinct shares, and, when a portable key was
// generated, I_ACCESS opens exactly one of them and yields the control (spec
// §62.1 rule 11).
func selfCheckInner(inner, control []byte, portable *age.X25519Identity) error {
	stanzas, err := agewrap.Stanzas(bytes.NewReader(inner))
	if err != nil {
		return fmt.Errorf("capsule: self-check: INNER_ACCESS_AGE: %w", err)
	}
	if err := agewrap.CheckAccessStanzas(stanzas, agewrap.AccessSlots); err != nil {
		return fmt.Errorf("capsule: self-check: %w", err)
	}
	if portable == nil {
		return nil
	}
	id, err := agewrap.NewAccessIdentity(agewrap.AccessSlots, portable)
	if err != nil {
		return err
	}
	got, err := decryptAll(inner, id)
	defer clear(got)
	if err != nil {
		return fmt.Errorf("capsule: self-check: the portable key does not open INNER_ACCESS_AGE: %w", err)
	}
	if !bytes.Equal(got, control) {
		return errors.New("capsule: self-check: INNER_ACCESS_AGE does not hold the control")
	}
	return nil
}

// selfCheckPayload checks the PAYLOAD_AGE just written: the plaintext handed
// to age was P bytes, so PAYLOAD_AGE has the length P gives, and I_PAYLOAD
// opens its header, whose MAC verifies (spec §62.1 rule 11). Without it an
// omitted padding would reveal the exact L, and the capsule would fail at
// step 17.
func selfCheckPayload(p *payloadWriter, payloadRaw []byte, padded uint64) error {
	if want := PayloadAgeLength(padded); p.n != want {
		return fmt.Errorf("capsule: self-check: PAYLOAD_AGE is %d bytes, P = %d gives %d", p.n, padded, want)
	}
	hdr, err := age.ExtractHeader(bytes.NewReader(p.head))
	if err != nil {
		return errors.New("capsule: self-check: the age header of PAYLOAD_AGE does not parse")
	}
	id, err := agewrap.NewPayloadIdentity(payloadRaw)
	if err != nil {
		return err
	}
	fileKey, err := age.DecryptHeader(hdr, id)
	clear(fileKey)
	if err != nil {
		return errors.New("capsule: self-check: I_PAYLOAD does not open the header of PAYLOAD_AGE")
	}
	return nil
}

// accessRecipients validates the policy options and returns the credentials
// of INNER_ACCESS_AGE, including R_ACCESS when a portable key is requested:
// from 1 to 16, X25519, canonical, not of low order, none twice (spec §37,
// §39, §62.1 rule 3).
func accessRecipients(opts EncryptOptions) ([]age.Recipient, *age.X25519Identity, error) {
	switch opts.Policy {
	case TimeOnly:
		if len(opts.Recipients) != 0 || opts.NewPortableKey {
			return nil, nil, errors.New("capsule: time_only takes no recipients and no portable key")
		}
		return nil, nil, nil
	case TimeAndKey:
	default:
		return nil, nil, fmt.Errorf("capsule: unknown access policy %d", opts.Policy)
	}
	n := len(opts.Recipients)
	if opts.NewPortableKey {
		n++
	}
	if n == 0 {
		return nil, nil, errors.New("capsule: time_and_key needs at least one recipient or a portable key")
	}
	if n > agewrap.AccessSlots {
		return nil, nil, fmt.Errorf("capsule: time_and_key takes at most %d credentials, recipients and portable key together; %d given", agewrap.AccessSlots, n)
	}
	var out []age.Recipient
	seen := make(map[string]bool)
	for i, r := range opts.Recipients {
		x, ok := r.(*age.X25519Recipient)
		if !ok || x == nil {
			return nil, nil, fmt.Errorf("capsule: recipient %d is %T; time_and_key accepts X25519 recipients only", i, r)
		}
		if err := agewrap.CheckX25519Recipient(x); err != nil {
			return nil, nil, fmt.Errorf("capsule: recipient %d: %w", i, err)
		}
		if seen[x.String()] {
			return nil, nil, fmt.Errorf("capsule: recipient %s listed twice; INNER_ACCESS_AGE holds one stanza per recipient", x)
		}
		seen[x.String()] = true
		out = append(out, x)
	}
	var portable *age.X25519Identity
	if opts.NewPortableKey {
		var err error
		if portable, err = age.GenerateX25519Identity(); err != nil {
			return nil, nil, err
		}
		out = append(out, portable.Recipient())
	}
	return out, portable, nil
}

// fillSlots returns the 16 recipients of INNER_ACCESS_AGE: the credentials,
// and in each slot left a dummy, the public key of a fresh X25519 identity
// whose private key is dropped at once and never stored or returned (spec
// §39), in a uniformly random order. age writes the stanzas in the order of
// its recipients, so this is the order of the stanzas.
func fillSlots(credentials []age.Recipient) ([]age.Recipient, error) {
	slots := append(make([]age.Recipient, 0, agewrap.AccessSlots), credentials...)
	for len(slots) < agewrap.AccessSlots {
		dummy, err := age.GenerateX25519Identity()
		if err != nil {
			return nil, err
		}
		slots = append(slots, dummy.Recipient())
	}
	return slots, permute(slots)
}

// permute puts s in a uniformly random order: Fisher-Yates with
// crypto/rand.Int, which draws without bias (spec §39).
func permute[T any](s []T) error {
	for i := len(s) - 1; i > 0; i-- {
		j, err := rand.Int(rand.Reader, big.NewInt(int64(i+1)))
		if err != nil {
			return err
		}
		k := int(j.Int64())
		s[i], s[k] = s[k], s[i]
	}
	return nil
}

// encryptAll produces a complete in-memory age file.
func encryptAll(plaintext []byte, recipients ...age.Recipient) ([]byte, error) {
	var buf bytes.Buffer
	w, err := age.Encrypt(&buf, recipients...)
	if err != nil {
		return nil, err
	}
	_, writeErr := w.Write(plaintext)
	if err := errors.Join(writeErr, w.Close()); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}