DateKeys/agewrap/agewrap_test.go

package agewrap_test

import (
	"bytes"
	"errors"
	"io"
	"strings"
	"testing"
	"time"

	"filippo.io/age"
	"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/agewrap"
	"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},
	}
	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},
	} {
		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)
	}
}

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(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(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)
	}
	stranger, _ := age.GenerateX25519Identity()
	accS, _ := agewrap.NewAccessIdentity(stranger)
	if _, err := decrypt(file, accS); !errors.Is(err, datekeys.ErrAccessInvalid) {
		t.Fatalf("stranger: %v", err)
	}
	if _, err := agewrap.NewAccessIdentity(); !errors.Is(err, datekeys.ErrAccessRequired) {
		t.Fatalf("no identity: %v", err)
	}
}

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(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])); !errors.Is(err, datekeys.ErrPolicyStructureMismatch) {
		t.Fatalf("repeated stanza: %v", err)
	}
	if err := agewrap.CheckAccessStanzas(nil); !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)
	})
}