package agewrap_test

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"reflect"
	"strings"
	"testing"
	"time"

	"filippo.io/age"
	"github.com/drand/drand/v2/crypto"
	"github.com/drand/kyber"
	bls "github.com/drand/kyber-bls12381"
	"github.com/drand/kyber/encrypt/ibe"
	"github.com/drand/tlock"

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

// tlockNetwork adapts the pinned profile to tlock.Network, to run the
// official tlock code path against our stanzas.
type tlockNetwork struct {
	p       *profile.Profile
	release provider.Release
}

func (n tlockNetwork) ChainHash() string            { return n.p.ChainHashHex() }
func (n tlockNetwork) Current(time.Time) uint64     { return 1 << 40 }
func (n tlockNetwork) SwitchChainHash(string) error { return errors.New("forbidden") }
func (n tlockNetwork) Scheme() crypto.Scheme {
	s, _ := n.p.DrandScheme()
	return *s
}
func (n tlockNetwork) PublicKey() kyber.Point {
	s, _ := n.p.DrandScheme()
	k := s.KeyGroup.Point()
	_ = k.UnmarshalBinary(n.p.PublicKey)
	return k
}
func (n tlockNetwork) Signature(round uint64) ([]byte, error) {
	if round != n.release.Round {
		return nil, errors.New("unknown round")
	}
	return n.release.Signature, nil
}

func encrypt(t *testing.T, plaintext []byte, r ...age.Recipient) []byte {
	t.Helper()
	var b bytes.Buffer
	w, err := age.Encrypt(&b, r...)
	if err != nil {
		t.Fatal(err)
	}
	w.Write(plaintext)
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
	return b.Bytes()
}

func decrypt(file []byte, id age.Identity) ([]byte, error) {
	r, err := age.Decrypt(bytes.NewReader(file), id)
	if err != nil {
		return nil, err
	}
	return io.ReadAll(r)
}

func TestTimeRecipientStanzaAndRoundTrip(t *testing.T) {
	p := profile.Quicknet()
	rec, err := agewrap.NewTimeRecipient(p, 1000)
	if err != nil {
		t.Fatal(err)
	}
	file := encrypt(t, []byte("control"), rec)
	st, err := agewrap.Stanzas(bytes.NewReader(file))
	if err != nil {
		t.Fatal(err)
	}
	if len(st) != 1 || st[0].Type != "tlock" || len(st[0].Args) != 2 || st[0].Args[0] != "1000" || st[0].Args[1] != p.ChainHashHex() || len(st[0].Body) != 128 {
		t.Fatalf("stanza %+v", st)
	}
	id, _ := agewrap.NewTimeIdentity(p, 1000, testkit.Release(1000))
	if got, err := decrypt(file, id); err != nil || string(got) != "control" {
		t.Fatalf("round trip: %q %v", got, err)
	}
}

// The stanza is the one of the tlock library and the tle CLI, in both
// directions (spec §32, plan §3.3).
func TestInteroperabilityWithTlockLibrary(t *testing.T) {
	p := profile.Quicknet()
	net := tlockNetwork{p: p, release: testkit.Release(1000)}

	rec, _ := agewrap.NewTimeRecipient(p, 1000)
	ours := encrypt(t, []byte("from datekeys"), rec)
	var out bytes.Buffer
	if err := tlock.New(net).Strict().Decrypt(&out, bytes.NewReader(ours)); err != nil || out.String() != "from datekeys" {
		t.Fatalf("tlock cannot open our file: %v", err)
	}

	var theirs bytes.Buffer
	if err := tlock.New(net).Strict().Encrypt(&theirs, strings.NewReader("from tlock"), 1000); err != nil {
		t.Fatal(err)
	}
	id, _ := agewrap.NewTimeIdentity(p, 1000, testkit.Release(1000))
	if got, err := decrypt(theirs.Bytes(), id); err != nil || string(got) != "from tlock" {
		t.Fatalf("we cannot open a tlock file: %v", err)
	}
}

func TestTimeRecipientCannotBeMixed(t *testing.T) {
	p := profile.Quicknet()
	a, _ := agewrap.NewTimeRecipient(p, 1000)
	b, _ := agewrap.NewTimeRecipient(p, 1000)
	x, _ := age.GenerateX25519Identity()
	for _, rs := range [][]age.Recipient{{a, x.Recipient()}, {a, b}} {
		if _, err := age.Encrypt(io.Discard, rs...); err == nil {
			t.Fatal("tlock recipient mixed with another recipient")
		}
	}
}

// Every rule is enforced in Unwrap even when the header MAC is valid, which
// is what a malicious creator produces (spec §27, §63).
func TestTimeIdentityStrictness(t *testing.T) {
	p := profile.Quicknet()
	rec, _ := agewrap.NewTimeRecipient(p, 1000)
	file, fk, err := testkit.Encrypt([]byte("control"), rec)
	if err != nil {
		t.Fatal(err)
	}
	rewrite := func(edit func([]*age.Stanza) []*age.Stanza) []byte {
		out, err := testkit.RewriteAge(file, fk, edit)
		if err != nil {
			t.Fatal(err)
		}
		// The rewritten header is authentic for age: the injected file key opens it.
		if _, err := decrypt(out, age.NewInjectedFileKeyIdentity(fk)); err != nil {
			t.Fatalf("rewritten file is not a valid age file: %v", err)
		}
		return out
	}
	backdoor, backdoorID, _ := testkit.X25519Stanza(fk)
	cases := []struct {
		name string
		file []byte
		want error
	}{
		{"extra X25519 stanza", rewrite(func(s []*age.Stanza) []*age.Stanza { return append(s, backdoor) }), datekeys.ErrPolicyStructureMismatch},
		{"stanza type changed", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Type = "tlock2"; return s }), datekeys.ErrPolicyStructureMismatch},
		{"third argument", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Args = append(s[0].Args, "x"); return s }), datekeys.ErrPolicyStructureMismatch},
		{"other round", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Args[0] = "1001"; return s }), datekeys.ErrRoundMismatch},
		{"round with leading zero", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Args[0] = "01000"; return s }), datekeys.ErrRoundMismatch},
		{"other chain hash", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Args[1] = strings.Repeat("0", 64); return s }), datekeys.ErrProfileMismatch},
		{"uppercase chain hash", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Args[1] = strings.ToUpper(s[0].Args[1]); return s }), datekeys.ErrProfileMismatch},
		{"corrupted tlock body", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body[100] ^= 1; return s }), datekeys.ErrIntegrity},
		{"truncated tlock body", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = s[0].Body[:127]; return s }), datekeys.ErrIntegrity},
		// Spec §12.2, §63 step 11: U || V || W with |U| = 96, and U the
		// canonical encoding of a point of G2 other than the point at
		// infinity. For a decoder that reduces c0 modulo p, c0 + p is U.
		{"tlock body of 129 bytes", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = append(s[0].Body, 0); return s }), datekeys.ErrIntegrity},
		{"U re-encoded with c0 + p", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = editU(t, s[0].Body, c0PlusP); return s }), datekeys.ErrIntegrity},
		{"U the point at infinity", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = editU(t, s[0].Body, infinity); return s }), datekeys.ErrIntegrity},
		{"U with the infinity flag and a payload", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = editU(t, s[0].Body, infinityWithPayload); return s }), datekeys.ErrIntegrity},
		{"U negated", rewrite(func(s []*age.Stanza) []*age.Stanza { s[0].Body = editU(t, s[0].Body, negated); return s }), datekeys.ErrIntegrity},
	}
	id, _ := agewrap.NewTimeIdentity(p, 1000, testkit.Release(1000))
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := decrypt(tc.file, id); !errors.Is(err, tc.want) {
				t.Fatalf("got %v, want %v", err, tc.want)
			}
		})
	}
	// Without the DateKeys rule, the backdoor stanza would open the file.
	if got, err := decrypt(cases[0].file, backdoorID); err != nil || string(got) != "control" {
		t.Fatalf("backdoor model broken: %v", err)
	}
}

func TestTimeIdentityRelease(t *testing.T) {
	p := profile.Quicknet()
	rec, _ := agewrap.NewTimeRecipient(p, 1000)
	file := encrypt(t, []byte("control"), rec)
	for _, tc := range []struct {
		name string
		rel  provider.Release
		want error
	}{
		{"release of another round", testkit.Release(1001), datekeys.ErrRoundMismatch},
		{"relabelled signature", provider.Release{Round: 1000, Signature: testkit.Release(1001).Signature}, datekeys.ErrReleaseInvalid},
		{"empty signature", provider.Release{Round: 1000}, datekeys.ErrReleaseInvalid},
		{"negated signature", provider.Release{Round: 1000, Signature: testkit.Negated(testkit.Release(1000).Signature)}, datekeys.ErrReleaseInvalid},
		{"signature the point at infinity", provider.Release{Round: 1000, Signature: testkit.Infinity(48)}, datekeys.ErrReleaseInvalid},
	} {
		id, _ := agewrap.NewTimeIdentity(p, 1000, tc.rel)
		if _, err := decrypt(file, id); !errors.Is(err, tc.want) {
			t.Errorf("%s: got %v, want %v", tc.name, err, tc.want)
		}
	}
	// An identity for another round refuses the stanza before using the release.
	id, _ := agewrap.NewTimeIdentity(p, 1001, testkit.Release(1001))
	if _, err := decrypt(file, id); !errors.Is(err, datekeys.ErrRoundMismatch) {
		t.Fatalf("identity for round 1001: %v", err)
	}
	// Spec §12.2: the published signature of a round re-encoded with x + p
	// is refused, although it is the same point for a decoder that reduces
	// x modulo p; the canonical one opens the file.
	rec, _ = agewrap.NewTimeRecipient(p, testkit.XPlusPRound)
	file = encrypt(t, []byte("control"), rec)
	canonical := testkit.Release(testkit.XPlusPRound)
	xPlusP, err := testkit.AddModulus(canonical.Signature, 0)
	if err != nil {
		t.Fatal(err)
	}
	id, _ = agewrap.NewTimeIdentity(p, testkit.XPlusPRound, provider.Release{Round: testkit.XPlusPRound, Signature: xPlusP})
	if _, err := decrypt(file, id); !errors.Is(err, datekeys.ErrReleaseInvalid) {
		t.Errorf("signature re-encoded with x + p: %v", err)
	}
	id, _ = agewrap.NewTimeIdentity(p, testkit.XPlusPRound, canonical)
	if got, err := decrypt(file, id); err != nil || string(got) != "control" {
		t.Errorf("canonical signature of round %d: %q %v", testkit.XPlusPRound, got, err)
	}
}

// Spec §63 step 11: H2 hashes the element of GT in the order of
// kilic/bls12-381, c1 before c0 at every level of the tower. The frozen
// vector H2(e(G1, G2)) of testdata/vectors/tlock_ibe.json pins the pairing and
// that serialization, and step 11 computed with them, sigma = V XOR
// H2(e(signature, U)) and FK_TIME = W XOR H4(sigma), recovers the file key
// that tlock unwraps. The reverse order, c0 first at every level as in the
// Fp12.toBytes of noble, gives another H2 and another key.
func TestTlockH2Vector(t *testing.T) {
	var golden testkit.IBEVectorFile
	if err := testkit.ReadJSON("../testdata/vectors/tlock_ibe.json", &golden); err != nil {
		t.Fatal(err)
	}
	if got, err := testkit.IBEVectors(); err != nil || !reflect.DeepEqual(got, golden) {
		t.Fatalf("testdata/vectors/tlock_ibe.json is stale: run genfixtures (%v)", err)
	}
	const (
		g1 = "97f1d3a73197d7942695638c4fa9ac0fc3688c4f9774b905a14e3a3f171bac586c55e83ff97a1aeffb3af00adb22c6bb"
		g2 = "93e02b6052719f607dacd3a088274f65596bd0d09920b61ab5da61bbdc7f5049334cf11213945d57e5ac7d055d042b7e" +
			"024aa2b2f08f0a91260805272dc51051c6e47ad4fa403b02b4510b647ae3d1770bac0326a805bbefd48056c8c121bdb8"
	)
	vec := golden.Vectors[0]
	if len(golden.Vectors) != 1 || vec.G1 != g1 || vec.G2 != g2 || vec.H2 != "cb87319f24560b5231579a09ad79f12e" {
		t.Fatalf("the frozen vector changed: %+v", golden.Vectors)
	}
	gt, err := hex.DecodeString(vec.GT)
	if err != nil || len(gt) != testkit.GTLen {
		t.Fatalf("gt of %d bytes: %v", len(gt), err)
	}
	if sum := sha256.Sum256(append([]byte("IBE-H2"), gt...)); hex.EncodeToString(sum[:testkit.H2Len]) != vec.H2 {
		t.Fatal("h2 is not SHA-256(\"IBE-H2\" || gt) truncated to 16 bytes")
	}
	if h := hex.EncodeToString(testkit.H2(reversed(gt))); h != "0118eea9d5971745f71e3c94926f1717" {
		t.Fatalf("H2 in the reverse order: %s", h)
	}

	// Step 11 on a stanza of round 1000 with the published release.
	p := profile.Quicknet()
	rec, _ := agewrap.NewTimeRecipient(p, 1000)
	stanzas, err := agewrap.Stanzas(bytes.NewReader(encrypt(t, []byte("control"), rec)))
	if err != nil {
		t.Fatal(err)
	}
	release := testkit.Release(1000)
	id, _ := agewrap.NewTimeIdentity(p, 1000, release)
	fileKey, err := id.Unwrap(stanzas)
	if err != nil {
		t.Fatal(err)
	}
	scheme, err := p.DrandScheme()
	if err != nil {
		t.Fatal(err)
	}
	body := stanzas[0].Body
	u, v, w := body[:len(body)-2*testkit.H2Len], body[len(body)-2*testkit.H2Len:len(body)-testkit.H2Len], body[len(body)-testkit.H2Len:]
	sig, point := scheme.SigGroup.Point(), scheme.KeyGroup.Point()
	if err := sig.UnmarshalBinary(release.Signature); err != nil {
		t.Fatal(err)
	}
	if err := point.UnmarshalBinary(u); err != nil {
		t.Fatal(err)
	}
	pairing, err := bls.NewBLS12381Suite().Pair(sig, point).MarshalBinary()
	if err != nil {
		t.Fatal(err)
	}
	for _, tc := range []struct {
		name  string
		gt    []byte
		opens bool
	}{
		{"the order of kilic/bls12-381", pairing, true},
		{"the reverse order", reversed(pairing), false},
	} {
		sigma := xor(v, testkit.H2(tc.gt))
		h4 := sha256.Sum256(append(ibe.H4Tag(), sigma...))
		if got := xor(w, h4[:testkit.H2Len]); bytes.Equal(got, fileKey) != tc.opens {
			t.Errorf("%s: FK_TIME %x, tlock unwraps %x", tc.name, got, fileKey)
		}
	}
}

// reversed returns the twelve 48-byte coordinates of a serialization of GT in
// reverse order: c0 before c1 at every level of the tower.
func reversed(gt []byte) []byte {
	out := make([]byte, 0, len(gt))
	for i := len(gt) - testkit.CoordinateLen; i >= 0; i -= testkit.CoordinateLen {
		out = append(out, gt[i:i+testkit.CoordinateLen]...)
	}
	return out
}

func xor(a, b []byte) []byte {
	out := make([]byte, len(a))
	for i := range a {
		out[i] = a[i] ^ b[i]
	}
	return out
}

// U edits of a Quicknet tlock stanza body U || V || W (spec §12.2, §63 step
// 11).
var (
	c0PlusP             = func(u []byte) ([]byte, error) { return testkit.AddModulus(u, testkit.CoordinateLen) }
	infinity            = func(u []byte) ([]byte, error) { return testkit.Infinity(len(u)), nil }
	infinityWithPayload = func(u []byte) ([]byte, error) { return testkit.InfinityWithPayload(u), nil }
	negated             = func(u []byte) ([]byte, error) { return testkit.Negated(u), nil }
)

func editU(t *testing.T, body []byte, edit func(u []byte) ([]byte, error)) []byte {
	t.Helper()
	out, err := testkit.EditU(edit)(body)
	if err != nil {
		t.Fatal(err)
	}
	return out
}

func TestPayloadIdentityStrictness(t *testing.T) {
	iPayload, _ := age.GenerateX25519Identity()
	raw, err := agewrap.RawX25519Identity(iPayload)
	if err != nil {
		t.Fatal(err)
	}
	id, err := agewrap.NewPayloadIdentity(raw)
	if err != nil {
		t.Fatal(err)
	}
	file, fk, _ := testkit.Encrypt([]byte("payload"), iPayload.Recipient())
	if got, err := decrypt(file, id); err != nil || string(got) != "payload" {
		t.Fatalf("round trip: %v", err)
	}
	backdoor, _, _ := testkit.X25519Stanza(fk)
	extra, _ := testkit.RewriteAge(file, fk, func(s []*age.Stanza) []*age.Stanza { return append(s, backdoor) })
	// Plain age accepts the file with I_PAYLOAD: the MAC is valid.
	if _, err := decrypt(extra, iPayload); err != nil {
		t.Fatalf("model broken: %v", err)
	}
	if _, err := decrypt(extra, id); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("extra stanza in PAYLOAD_AGE: %v", err)
	}
	other, _ := age.GenerateX25519Identity()
	if _, err := decrypt(encrypt(t, []byte("x"), other.Recipient()), id); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("payload of another control: %v", err)
	}
	pw, _ := age.NewScryptRecipient("password")
	pw.SetWorkFactor(10)
	if _, err := decrypt(encrypt(t, []byte("x"), pw), id); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("scrypt payload: %v", err)
	}
}

func TestAccessIdentityStrictness(t *testing.T) {
	a, _ := age.GenerateX25519Identity()
	b, _ := age.GenerateX25519Identity()
	file, fk, _ := testkit.Encrypt([]byte("control"), a.Recipient(), b.Recipient())
	for _, id := range []*age.X25519Identity{a, b} {
		acc, _ := agewrap.NewAccessIdentity(0, id)
		if got, err := decrypt(file, acc); err != nil || string(got) != "control" {
			t.Fatalf("recipient cannot open: %v", err)
		}
	}
	// A non-X25519 stanza is rejected although plain age would accept the file.
	odd := &age.Stanza{Type: "scrypt", Args: []string{"c2FsdHNhbHRzYWx0c2FsdA", "10"}, Body: make([]byte, 32)}
	withOdd, _ := testkit.RewriteAge(file, fk, func(s []*age.Stanza) []*age.Stanza { return append(s, odd) })
	if _, err := decrypt(withOdd, a); err != nil {
		t.Fatalf("model broken: %v", err)
	}
	acc, _ := agewrap.NewAccessIdentity(0, a)
	if _, err := decrypt(withOdd, acc); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("non-X25519 stanza: %v", err)
	}
	// Two stanzas for the same recipient.
	dup, _ := a.Recipient().Wrap(fk)
	withDup, _ := testkit.RewriteAge(file, fk, func(s []*age.Stanza) []*age.Stanza { return append(s, dup[0]) })
	if _, err := decrypt(withDup, acc); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("two stanzas for one recipient: %v", err)
	}
	// Spec §63 step 13, §69.1: the verdict does not depend on the order of
	// the identities. b unwraps exactly one stanza of withDup and a two, in
	// either order.
	for i, ids := range [][]age.Identity{{a, b}, {b, a}} {
		both, _ := agewrap.NewAccessIdentity(0, ids...)
		if _, err := decrypt(withDup, both); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
			t.Fatalf("two stanzas for a, order %d: %v", i, err)
		}
	}
	stranger, _ := age.GenerateX25519Identity()
	for _, ids := range [][]age.Identity{{stranger, b}, {b, stranger}} {
		mixed, _ := agewrap.NewAccessIdentity(0, ids...)
		if got, err := decrypt(file, mixed); err != nil || string(got) != "control" {
			t.Fatalf("a stranger and a recipient: %v", err)
		}
	}
	accS, _ := agewrap.NewAccessIdentity(0, stranger)
	if _, err := decrypt(file, accS); !errors.Is(err, datekeys.ErrAccessInvalid) {
		t.Fatalf("stranger: %v", err)
	}
	if _, err := agewrap.NewAccessIdentity(0); !errors.Is(err, datekeys.ErrAccessRequired) {
		t.Fatalf("no identity: %v", err)
	}
}

// Spec §39, §63 steps 12 and 13: in format 2 INNER_ACCESS_AGE holds exactly
// AccessSlots stanzas, a rule checked on its own and again by the identity;
// format 1 takes one or more.
func TestAccessSlots(t *testing.T) {
	a, _ := age.GenerateX25519Identity()
	for _, n := range []int{1, 15, 16, 17} {
		recipients := []age.Recipient{a.Recipient()}
		for len(recipients) < n {
			id, _ := age.GenerateX25519Identity()
			recipients = append(recipients, id.Recipient())
		}
		file := encrypt(t, []byte("control"), recipients...)
		st, err := agewrap.Stanzas(bytes.NewReader(file))
		if err != nil {
			t.Fatal(err)
		}
		acc, _ := agewrap.NewAccessIdentity(agewrap.AccessSlots, a)
		_, openErr := decrypt(file, acc)
		for _, err := range []error{agewrap.CheckAccessStanzas(st, agewrap.AccessSlots), openErr} {
			if (n == agewrap.AccessSlots) != (err == nil) || (err != nil && !errors.Is(err, datekeys.ErrPolicyStructureMismatch)) {
				t.Fatalf("%d stanzas in format 2: %v", n, err)
			}
		}
		acc1, _ := agewrap.NewAccessIdentity(0, a)
		if err := agewrap.CheckAccessStanzas(st, 0); err != nil {
			t.Fatalf("%d stanzas in format 1: %v", n, err)
		}
		if got, err := decrypt(file, acc1); err != nil || string(got) != "control" {
			t.Fatalf("%d stanzas in format 1: %v", n, err)
		}
	}
	if _, err := agewrap.NewAccessIdentity(-1, a); err == nil {
		t.Fatal("negative slots accepted")
	}
}

// Spec §37, §62.1 rule 3, §76 change 10: the X25519 recipients a writer must
// reject, non-canonical or of low order, whose stanza nobody or anybody
// opens; the rules of the reader cannot tell, because the stanza does not
// hold the recipient.
func TestNonCanonicalRecipients(t *testing.T) {
	x, _ := age.GenerateX25519Identity()
	if err := agewrap.CheckX25519Recipient(x.Recipient()); err != nil {
		t.Fatalf("a generated recipient: %v", err)
	}
	fromRaw := func(b []byte) *age.X25519Recipient {
		t.Helper()
		s, err := bech32.Encode("age", b)
		if err != nil {
			t.Fatal(err)
		}
		r, err := age.ParseX25519Recipient(s)
		if err != nil {
			t.Fatal(err)
		}
		return r
	}
	ff := func(first, last byte) []byte {
		return append(append([]byte{first}, bytes.Repeat([]byte{0xff}, 30)...), last)
	}
	raw, _ := agewrap.RawX25519Recipient(x.Recipient())
	high := bytes.Clone(raw)
	high[31] |= 0x80
	mustHex := func(s string) []byte {
		b, err := hex.DecodeString(s)
		if err != nil {
			t.Fatal(err)
		}
		return b
	}
	for name, u := range map[string][]byte{
		"bit 255 set":     high,
		"u = p":           ff(0xed, 0x7f),
		"u = p + 1":       ff(0xee, 0x7f),
		"u = 2^255 - 1":   ff(0xff, 0x7f),
		"low order, 0":    make([]byte, 32),
		"low order, 1":    append([]byte{1}, make([]byte, 31)...),
		"low order, p-1":  ff(0xec, 0x7f),
		"low order, 8 #1": mustHex("e0eb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b800"),
		"low order, 8 #2": mustHex("5f9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f1157"),
	} {
		if err := agewrap.CheckX25519Recipient(fromRaw(u)); err == nil {
			t.Errorf("%s accepted", name)
		}
	}
	// The cases of spec §76: age encrypts to the recipient with bit 255 set
	// a stanza that the identity cannot open, and refuses to encrypt to the
	// zero point.
	if _, err := decrypt(encrypt(t, []byte("x"), fromRaw(high)), x); err == nil {
		t.Fatal("the identity opens the stanza of its non-canonical recipient")
	}
	w, err := age.Encrypt(io.Discard, fromRaw(make([]byte, 32)))
	if err == nil {
		err = w.Close()
	}
	if err == nil {
		t.Fatal("age encrypts to a point of low order")
	}
}

func TestStanzasProbe(t *testing.T) {
	if _, err := agewrap.Stanzas(strings.NewReader("not age")); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("garbage: %v", err)
	}
	x, _ := age.GenerateX25519Identity()
	file := encrypt(t, []byte("x"), x.Recipient())
	st, err := agewrap.Stanzas(bytes.NewReader(file))
	if err != nil || len(st) != 1 || st[0].Type != "X25519" {
		t.Fatalf("%+v %v", st, err)
	}
	// Probing never needs a secret and leaves the stanzas untouched for age.
	if _, err := decrypt(file, x); err != nil {
		t.Fatal(err)
	}
}

func TestRawKeys(t *testing.T) {
	id, _ := age.GenerateX25519Identity()
	raw, err := agewrap.RawX25519Identity(id)
	if err != nil || len(raw) != 32 {
		t.Fatal(err)
	}
	back, err := agewrap.X25519IdentityFromRaw(raw)
	if err != nil || back.String() != id.String() {
		t.Fatal("identity round trip")
	}
	pub, err := agewrap.RawX25519Recipient(id.Recipient())
	if err != nil || len(pub) != 32 {
		t.Fatal(err)
	}
	if _, err := agewrap.X25519IdentityFromRaw(raw[:31]); err == nil {
		t.Fatal("31-byte identity accepted")
	}
}

func TestConstructorsRejectInvalidInput(t *testing.T) {
	p := profile.Quicknet()
	for _, round := range []uint64{0, p.MaxRound() + 1} {
		if _, err := agewrap.NewTimeRecipient(p, round); !errors.Is(err, datekeys.ErrDateKeyInvalid) {
			t.Errorf("round %d: %v", round, err)
		}
	}
	for name, edit := range map[string]func(p *profile.Profile){
		"unknown scheme":         func(p *profile.Profile) { p.Scheme = "nope" },
		"public key not a point": func(p *profile.Profile) { p.PublicKey = bytes.Repeat([]byte{0xff}, 96) },
		"identity element":       func(p *profile.Profile) { p.PublicKey = append([]byte{0xc0}, make([]byte, 95)...) },
	} {
		bad := profile.Quicknet()
		edit(bad)
		if _, err := agewrap.NewTimeRecipient(bad, 1000); !errors.Is(err, datekeys.ErrUnknownProfile) {
			t.Errorf("recipient, %s: %v", name, err)
		}
		if _, err := agewrap.NewTimeIdentity(bad, 1000, testkit.Release(1000)); !errors.Is(err, datekeys.ErrUnknownProfile) {
			t.Errorf("identity, %s: %v", name, err)
		}
	}
	if _, err := agewrap.NewPayloadIdentity(make([]byte, 31)); err == nil {
		t.Fatal("31-byte I_PAYLOAD accepted")
	}
}

func TestMalformedX25519Stanzas(t *testing.T) {
	x, _ := age.GenerateX25519Identity()
	file, fk, _ := testkit.Encrypt([]byte("data"), x.Recipient())
	malformed, err := testkit.RewriteAge(file, fk, func(s []*age.Stanza) []*age.Stanza {
		s[0].Args = append(s[0].Args, "extra")
		return s
	})
	if err != nil {
		t.Fatal(err)
	}
	raw, _ := agewrap.RawX25519Identity(x)
	pid, _ := agewrap.NewPayloadIdentity(raw)
	if _, err := decrypt(malformed, pid); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("payload: %v", err)
	}
	acc, _ := agewrap.NewAccessIdentity(0, x)
	if _, err := decrypt(malformed, acc); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("access: %v", err)
	}
	st, _ := agewrap.Stanzas(bytes.NewReader(file))
	if err := agewrap.CheckAccessStanzas(append(st, st[0]), 0); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("repeated stanza: %v", err)
	}
	if err := agewrap.CheckAccessStanzas(nil, 0); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("no stanza: %v", err)
	}
	if err := agewrap.CheckPayloadStanzas(append(st, st[0])); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("two payload stanzas: %v", err)
	}
}

func FuzzStanzas(f *testing.F) {
	x, _ := age.GenerateX25519Identity()
	var b bytes.Buffer
	w, _ := age.Encrypt(&b, x.Recipient())
	w.Close()
	f.Add(b.Bytes())
	f.Add([]byte("age-encryption.org/v1\n-> X25519 a\n\n--- AAAA\n"))
	f.Fuzz(func(t *testing.T, in []byte) {
		st, err := agewrap.Stanzas(bytes.NewReader(in))
		if err != nil && !errors.Is(err, datekeys.ErrIntegrity) {
			t.Fatalf("unexpected error class: %v", err)
		}
		_ = agewrap.CheckPayloadStanzas(st)
		_ = agewrap.CheckAccessStanzas(st, 0)
		_ = agewrap.CheckAccessStanzas(st, agewrap.AccessSlots)
	})
}