pqsafe/pkg/envelope/envelope_test.go

// Copyright (C) 2025-2026, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package envelope_test

import (
	"bytes"
	"crypto/rand"
	"crypto/sha512"
	"errors"
	"testing"
	"time"

	"github.com/luxfi/zwing"

	"github.com/luxfi/pqsafe/pkg/ctx"
	"github.com/luxfi/pqsafe/pkg/envelope"
	"github.com/luxfi/pqsafe/pkg/manifest"
)

// E2E round-trip via zwing identities: sender seals a file for two
// recipients; each opens and recovers byte-identical plaintext.
func TestSealOpenRoundTrip(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)
	bob := mustGen(t)

	plaintext := make([]byte, 200_000) // crosses several chunks at 64 KiB
	if _, err := rand.Read(plaintext); err != nil {
		t.Fatalf("rand.Read: %v", err)
	}

	sealed, err := envelope.Seal(envelope.SealInput{
		Plaintext: plaintext,
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: sender.Public()}, // placeholder; real recipients below
		},
	})
	_ = sealed
	_ = err

	// Real call with distinct recipients:
	sealed, err = envelope.Seal(envelope.SealInput{
		Plaintext: plaintext,
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
			{ID: "bob@gov", Public: bob.Public()},
		},
		PolicyLabel: "CUI",
		TTL:         24 * time.Hour,
	})
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}
	if len(sealed.Capsules) != 2 {
		t.Fatalf("expected 2 capsules, got %d", len(sealed.Capsules))
	}

	for _, r := range []struct {
		who string
		id  string
		key *zwing.Identity
	}{
		{"alice", "alice@mil", alice},
		{"bob", "bob@gov", bob},
	} {
		got, err := envelope.Open(envelope.OpenInput{
			ManifestBytes: sealed.ManifestBytes,
			ManifestSig:   sealed.ManifestSig,
			Chunks:        sealed.Chunks,
			Capsules:      sealed.Capsules,
			Recipient:     r.key,
			RecipientID:   r.id,
			Now:           time.Now(),
		})
		if err != nil {
			t.Fatalf("Open (%s): %v", r.who, err)
		}
		if !bytes.Equal(got, plaintext) {
			t.Fatalf("%s plaintext mismatch (len got=%d want=%d)", r.who, len(got), len(plaintext))
		}
	}
}

func TestWrongRecipientCannotOpen(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)
	bob := mustGen(t)
	mallory := mustGen(t)

	sealed, _ := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("payload"), Sender: sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
			{ID: "bob@gov", Public: bob.Public()},
		},
	})

	_, err := envelope.Open(envelope.OpenInput{
		ManifestBytes: sealed.ManifestBytes,
		ManifestSig:   sealed.ManifestSig,
		Chunks:        sealed.Chunks,
		Capsules:      sealed.Capsules,
		Recipient:     mallory,
		RecipientID:   "alice@mil",
	})
	if err == nil {
		t.Fatal("expected Open to fail for wrong recipient key")
	}
}

func TestCorruptChunkFails(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	sealed, _ := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("hello-world-this-is-the-payload"),
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
		},
	})

	sealed.Chunks[0][0] ^= 0xFF

	_, err := envelope.Open(envelope.OpenInput{
		ManifestBytes: sealed.ManifestBytes,
		ManifestSig:   sealed.ManifestSig,
		Chunks:        sealed.Chunks,
		Capsules:      sealed.Capsules,
		Recipient:     alice,
		RecipientID:   "alice@mil",
	})
	if err == nil {
		t.Fatal("expected Open to fail on corrupted chunk")
	}
}

func TestTamperedManifestFails(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	sealed, _ := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("hi"),
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
		},
	})

	sealed.ManifestBytes[10] ^= 0x01

	_, err := envelope.Open(envelope.OpenInput{
		ManifestBytes: sealed.ManifestBytes,
		ManifestSig:   sealed.ManifestSig,
		Chunks:        sealed.Chunks,
		Capsules:      sealed.Capsules,
		Recipient:     alice,
		RecipientID:   "alice@mil",
	})
	if err == nil {
		t.Fatal("expected Open to fail on tampered manifest")
	}
}

func TestSwapCapsulesFails(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)
	bob := mustGen(t)

	sealed, _ := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("classified-payload"),
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
			{ID: "bob@gov", Public: bob.Public()},
		},
	})

	sealed.Capsules[0], sealed.Capsules[1] = sealed.Capsules[1], sealed.Capsules[0]

	_, err := envelope.Open(envelope.OpenInput{
		ManifestBytes: sealed.ManifestBytes,
		ManifestSig:   sealed.ManifestSig,
		Chunks:        sealed.Chunks,
		Capsules:      sealed.Capsules,
		Recipient:     alice,
		RecipientID:   "alice@mil",
	})
	if err == nil {
		t.Fatal("expected Open to fail when capsules are swapped")
	}
}

func TestExpiredEnvelopeRefused(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	sealed, _ := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("x"),
		Sender:    sender,
		Recipients: []envelope.Recipient{
			{ID: "alice@mil", Public: alice.Public()},
		},
		TTL: 1 * time.Hour,
	})

	_, err := envelope.Open(envelope.OpenInput{
		ManifestBytes: sealed.ManifestBytes,
		ManifestSig:   sealed.ManifestSig,
		Chunks:        sealed.Chunks,
		Capsules:      sealed.Capsules,
		Recipient:     alice,
		RecipientID:   "alice@mil",
		Now:           time.Now().Add(48 * time.Hour),
	})
	if err == nil {
		t.Fatal("expected expired envelope to be refused")
	}
}

func mustGen(t *testing.T) *zwing.Identity {
	t.Helper()
	id, err := zwing.GenerateIdentity()
	if err != nil {
		t.Fatalf("zwing.GenerateIdentity: %v", err)
	}
	return id
}

// reseal rewrites the manifest after the caller mutates it, re-signs it
// with the same sender, and returns the new manifest bytes + signature
// so Open can be exercised against attacker-controlled manifest fields.
func reseal(t *testing.T, sender *zwing.Identity, mf *manifest.Manifest) ([]byte, []byte) {
	t.Helper()
	mb, err := mf.MarshalCanonical()
	if err != nil {
		t.Fatalf("MarshalCanonical: %v", err)
	}
	h, err := mf.Hash()
	if err != nil {
		t.Fatalf("Hash: %v", err)
	}
	sig := sender.Sign([]byte(ctx.Manifest), h)
	return mb, sig
}

// A malicious sender (or downgrade attacker who has the signing key)
// tries to embed a non-canonical algorithm string. Open MUST refuse
// before any cryptographic operation runs.
func TestAlgorithmAllowListRefusesDowngrade(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	cases := []struct {
		name  string
		mut   func(*manifest.Manifest)
		errIs error
	}{
		{"hash", func(m *manifest.Manifest) { m.HashAlgo = "SHA-256" }, envelope.ErrBadAlgorithm},
		{"kdf", func(m *manifest.Manifest) { m.KDF = "HKDF-SHA-256" }, envelope.ErrBadAlgorithm},
		{"kem", func(m *manifest.Manifest) { m.KEM = "X25519" }, envelope.ErrBadAlgorithm},
		{"sig_classical", func(m *manifest.Manifest) { m.SigClassical = "ECDSA-P256" }, envelope.ErrBadAlgorithm},
		{"sig_pq", func(m *manifest.Manifest) { m.SigPQ = "Falcon-512" }, envelope.ErrBadAlgorithm},
		{"chunk_size", func(m *manifest.Manifest) { m.ChunkSize = 4096 }, envelope.ErrChunkSize},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			sealed, err := envelope.Seal(envelope.SealInput{
				Plaintext: []byte("payload"), Sender: sender,
				Recipients: []envelope.Recipient{{ID: "alice@mil", Public: alice.Public()}},
			})
			if err != nil {
				t.Fatalf("Seal: %v", err)
			}

			mf, err := manifest.Unmarshal(sealed.ManifestBytes)
			if err != nil {
				t.Fatalf("Unmarshal: %v", err)
			}
			tc.mut(mf)
			mb, sig := reseal(t, sender, mf)

			_, err = envelope.Open(envelope.OpenInput{
				ManifestBytes: mb, ManifestSig: sig,
				Chunks: sealed.Chunks, Capsules: sealed.Capsules,
				Recipient: alice, RecipientID: "alice@mil",
			})
			if !errors.Is(err, tc.errIs) {
				t.Fatalf("expected %v, got %v", tc.errIs, err)
			}
		})
	}
}

// Sealed.EnvelopeID is the SHA-384 hash of the canonical manifest. It
// is a 48-byte authenticated identifier callers can use as an external
// dedupe / audit key.
func TestEnvelopeIDIsManifestHash(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	sealed, err := envelope.Seal(envelope.SealInput{
		Plaintext: []byte("payload"), Sender: sender,
		Recipients: []envelope.Recipient{{ID: "alice@mil", Public: alice.Public()}},
	})
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	mf, err := manifest.Unmarshal(sealed.ManifestBytes)
	if err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}
	wantHash, err := mf.Hash()
	if err != nil {
		t.Fatalf("manifest.Hash: %v", err)
	}
	if !bytes.Equal(sealed.EnvelopeID, wantHash) {
		t.Fatal("Sealed.EnvelopeID must equal SHA-384(canonical(manifest))")
	}
	if len(sealed.EnvelopeID) != 48 {
		t.Fatalf("envelope_id must be 48 bytes (SHA-384), got %d", len(sealed.EnvelopeID))
	}
}

// envelope_id changes when ANY content-bearing manifest field is altered
// — including MerkleRoot, PolicyLabel, PolicyURI, TTL, BulkAEAD, and
// ChunkSize. This is the forgery-resistance property Red HIGH #3 demanded:
// a sender holding signing keys can NOT produce two envelopes with
// identical envelope_id but different content. The signature is over the
// manifest hash, so any forged manifest with the same envelope_id is
// either byte-identical to the original (no forgery) or fails sig
// verification (caught before any decryption).
func TestEnvelopeIDChangesWithContentFields(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	mkSealed := func() *envelope.Sealed {
		s, err := envelope.Seal(envelope.SealInput{
			Plaintext:   []byte("payload"),
			Sender:      sender,
			Recipients:  []envelope.Recipient{{ID: "alice@mil", Public: alice.Public()}},
			PolicyLabel: "CUI",
			PolicyURI:   "https://example.test/cui",
			TTL:         1 * time.Hour,
		})
		if err != nil {
			t.Fatalf("Seal: %v", err)
		}
		return s
	}
	baseSealed := mkSealed()
	baseMf, err := manifest.Unmarshal(baseSealed.ManifestBytes)
	if err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}
	baseID := baseSealed.EnvelopeID

	mutate := func(name string, m func(*manifest.Manifest)) {
		t.Run(name, func(t *testing.T) {
			mf, err := manifest.Unmarshal(baseSealed.ManifestBytes)
			if err != nil {
				t.Fatalf("Unmarshal: %v", err)
			}
			m(mf)
			mb, err := mf.MarshalCanonical()
			if err != nil {
				t.Fatalf("MarshalCanonical: %v", err)
			}
			h := sha512.Sum384(mb)
			if bytes.Equal(h[:], baseID) {
				t.Fatalf("envelope_id unchanged after mutating %q — content field not in transcript", name)
			}
		})
	}

	mutate("MerkleRoot", func(m *manifest.Manifest) { m.MerkleRoot[0] ^= 0xFF })
	mutate("PolicyLabel", func(m *manifest.Manifest) { m.PolicyLabel = "TOP-SECRET" })
	mutate("PolicyURI", func(m *manifest.Manifest) { m.PolicyURI = "https://other.test/policy" })
	mutate("TTL", func(m *manifest.Manifest) { m.TTL = (48 * time.Hour).String() })
	mutate("BulkAEAD", func(m *manifest.Manifest) { m.BulkAEAD = "ChaCha20-Poly1305" })
	mutate("ChunkSize", func(m *manifest.Manifest) { m.ChunkSize = 4096 })
	mutate("Chunks", func(m *manifest.Manifest) {
		// Flip a byte in the first chunk's SHA — content-bearing because
		// chunk SHAs commit to ciphertexts.
		m.Chunks[0].SHA[0] ^= 0xFF
	})
	mutate("Created", func(m *manifest.Manifest) {
		// Created is in the AAD transcript AND in the manifest hash, so
		// it must change envelope_id too.
		m.Created = time.Now().UTC().Add(time.Hour).Format(time.RFC3339)
	})

	_ = baseMf
}

// A re-signed manifest with any altered content-bearing field must fail
// signature verification, because the signature is over the manifest
// hash and that hash IS the envelope_id. This closes the "same envelope
// _id, different content" forgery class end-to-end through Open.
func TestForgedContentRefusedAtSignatureVerify(t *testing.T) {
	sender := mustGen(t)
	alice := mustGen(t)

	sealed, err := envelope.Seal(envelope.SealInput{
		Plaintext:   []byte("payload"),
		Sender:      sender,
		Recipients:  []envelope.Recipient{{ID: "alice@mil", Public: alice.Public()}},
		PolicyLabel: "CUI",
	})
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	mf, err := manifest.Unmarshal(sealed.ManifestBytes)
	if err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}
	// Forge: change PolicyLabel without re-signing. Manifest hash now
	// differs from what's signed; Open must refuse at signature verify.
	mf.PolicyLabel = "PUBLIC"
	mb, err := mf.MarshalCanonical()
	if err != nil {
		t.Fatalf("MarshalCanonical: %v", err)
	}

	_, err = envelope.Open(envelope.OpenInput{
		ManifestBytes: mb, ManifestSig: sealed.ManifestSig,
		Chunks: sealed.Chunks, Capsules: sealed.Capsules,
		Recipient: alice, RecipientID: "alice@mil",
	})
	if !errors.Is(err, envelope.ErrBadSig) {
		t.Fatalf("expected ErrBadSig, got %v", err)
	}
}