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 EncryptFiles and 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: no
	// existing one is accepted. The recipients and the portable key are the
	// credentials of the capsule: from 1 to 16 (spec §39).
	NewPortableKey bool
	// Length is L for Encrypt, 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. EncryptFiles
	// computes L, the length of BODY, from the files, and requires 0.
	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
	// Comment and Author are the comment and the declared author of the
	// head that EncryptFiles writes, "" when absent (spec §29.4, §29.6):
	// the comment of 1 to 16384 bytes, in which EncryptFiles turns CR LF, and
	// a lone CR, into LF, and the declared author of 1 to 256. The declared
	// author is text of the creator and proves nothing (spec §55.1).
	Comment, Author string
	// HeadCritical and HeadNoncritical are the extensions of the head that
	// EncryptFiles writes, sealed in PAYLOAD_AGE (spec §29.4, §54).
	HeadCritical, HeadNoncritical []extension.Extension
	// TestVectors lets Encrypt write format 2, which only a generator of
	// test vectors may write (spec §62.1 rule 1, §70). EncryptFiles, which
	// writes format 3, ignores it.
	TestVectors bool
	// 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 EncryptFiles or 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: Format3 by EncryptFiles, Format2 by
	// Encrypt. Length is L, the length of the content, BODY in format 3,
	// Padding the padding rule and PaddedLength P = rule(L), the length of
	// the plaintext of PAYLOAD_AGE (spec §29.1, §29.2).
	Format       Format
	Length       uint64
	Padding      Padding
	PaddedLength uint64
	// Head is the head that EncryptFiles wrote: the files in the byte order
	// of their paths, with their layout and SHA-256, and the comment as
	// written. Nil for Encrypt.
	Head *Head
	// 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 and
// §62 of v0.9). Only a generator of test vectors may write format 2 (spec
// §62.1 rule 1, §70): Encrypt fails unless opts.TestVectors is set, and
// takes no comment, author or head extensions, which format 2 has no place
// for. Capsules are written with EncryptFiles.
//
// 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) {
	switch {
	case !opts.TestVectors:
		return nil, errors.New("capsule: Encrypt writes format 2, which only a generator of test vectors may write (spec §62.1 rule 1): EncryptFiles writes format 3")
	case opts.Comment != "" || opts.Author != "" || opts.HeadCritical != nil || opts.HeadNoncritical != nil:
		return nil, errors.New("capsule: format 2 has no head: Comment, Author and the head extensions are for EncryptFiles")
	case opts.Length < 0:
		return nil, fmt.Errorf("capsule: EncryptOptions.Length %d is negative", opts.Length)
	}
	s, err := newSealer(opts, uint64(opts.Length))
	if err != nil {
		return nil, err
	}
	return s.write(dst, Format2, uint64(opts.Length), func(w io.Writer) error {
		return copyExactly(w, src, opts.Length)
	})
}

// sealer writes what the writers of both formats share: the steps of spec
// §61 and §62 other than those of the content.
type sealer struct {
	opts        EncryptOptions
	code        Padding
	dk          datekey.DateKey
	unlock      time.Time
	credentials []age.Recipient
	portable    *age.X25519Identity
}

// newSealer validates the options that do not depend on the content, with
// length, a first L, checked against its maximum, and resolves the DateKey
// locally (spec §15, §62.1 rules 2, 3 and 8).
func newSealer(opts EncryptOptions, length uint64) (*sealer, 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))
	}
	s := &sealer{opts: opts, code: opts.Padding}
	if s.code == 0 {
		s.code = Reforzado
	}
	if _, err := PaddedLength(length, s.code); err != nil {
		return nil, err
	}

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

// write writes a capsule of format f whose content, of length bytes, body
// writes into the plaintext of PAYLOAD_AGE; write adds the zeros of the
// padding up to P (spec §29.1).
func (s *sealer) write(dst io.Writer, f Format, length uint64, body func(w io.Writer) error) (*Result, error) {
	opts := s.opts
	padded, err := PaddedLength(length, s.code)
	if err != nil {
		return nil, err
	}
	var portableRaw []byte
	if s.portable != nil {
		if portableRaw, err = agewrap.RawX25519Identity(s.portable); err != nil {
			return nil, err
		}
		defer clear(portableRaw)
	}

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

	// Step 6: 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 7 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(s.credentials); err != nil {
			return nil, err
		}
	}

	// Step 7 of spec §61 (8 of §62): PUBLIC_HEADER.
	header := &Header{CapsuleID: capsuleID, DateKey: s.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(opts.Profile, s.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, s.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 8 to 11 of spec §61 (9 to 12 of §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 control of version 2 or
	// 3 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: s.code,
	}
	draft, err := EncodeControl(ctrl, f)
	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: f, 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, with L and the padding code.
	controlBytes, err := EncodeControl(ctrl, f)
	if err != nil {
		return nil, err
	}
	defer clear(controlBytes)
	if err := selfCheckControl(controlBytes, f); 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
	}
	content := &countingWriter{w: aw}
	if err := body(content); err != nil {
		return nil, err
	}
	if content.n != length {
		return nil, fmt.Errorf("capsule: internal error: %d bytes of content, L = %d", content.n, length)
	}
	if err := writeZeros(aw, padded-length); 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: s.dk, UnlockAt: s.unlock, CapsuleID: capsuleID, Format: f, Length: length, Padding: s.code, PaddedLength: padded}
	if s.portable != nil {
		// The portable identity as 32 raw bytes in a .dkk (§62 step 18).
		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
}

// copyExactly writes to w exactly length bytes of src. 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 copyExactly(w io.Writer, src io.Reader, length int64) error {
	n, err := io.CopyN(w, 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
	}
	return nil
}

// writeZeros writes n zeros to w: the padding of spec §29.1.
func writeZeros(w io.Writer, n uint64) error {
	zeros := make([]byte, min(n, 16<<10))
	for n > 0 {
		k := min(n, uint64(len(zeros)))
		if _, err := w.Write(zeros[:k]); err != nil {
			return err
		}
		n -= k
	}
	return nil
}

// countingWriter counts the bytes written to w.
type countingWriter struct {
	w io.Writer
	n uint64
}

func (c *countingWriter) Write(b []byte) (int, error) {
	n, err := c.w.Write(b)
	c.n += uint64(n)
	return n, err
}

// 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 of format f 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 (spec §62.1 rules 11 and 17).
func selfCheckControl(b []byte, f Format) error {
	c, err := DecodeControl(b, f)
	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
}