selfpost/internal/secretfile/secretfile_test.go

package secretfile

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"testing"
)

// encrypt is the round-trip helper: seal data and return the envelope bytes.
func encrypt(t *testing.T, typ PayloadType, password string, data []byte) []byte {
	t.Helper()
	var buf bytes.Buffer
	w, err := NewWriter(&buf, typ, password)
	if err != nil {
		t.Fatalf("NewWriter: %v", err)
	}
	if _, err := w.Write(data); err != nil {
		t.Fatalf("Write: %v", err)
	}
	if err := w.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	return buf.Bytes()
}

func decrypt(t *testing.T, enc []byte, password string) ([]byte, PayloadType, error) {
	t.Helper()
	r, err := NewReader(bytes.NewReader(enc), password)
	if err != nil {
		return nil, 0, err
	}
	out, err := io.ReadAll(r)
	return out, r.Type(), err
}

func TestRoundTripSizes(t *testing.T) {
	// Empty, sub-chunk, exactly one chunk, and several chunks with a partial
	// tail: the boundaries where chunk framing tends to break.
	sizes := []int{0, 1, 1000, chunkSize - 1, chunkSize, chunkSize + 1, 3*chunkSize + 77}
	for _, n := range sizes {
		data := make([]byte, n)
		for i := range data {
			data[i] = byte(i * 7)
		}
		enc := encrypt(t, TypeFullBackup, "correct horse battery staple", data)
		got, typ, err := decrypt(t, enc, "correct horse battery staple")
		if err != nil {
			t.Fatalf("size %d: decrypt: %v", n, err)
		}
		if typ != TypeFullBackup {
			t.Errorf("size %d: type = %v, want full backup", n, typ)
		}
		if !bytes.Equal(got, data) {
			t.Errorf("size %d: plaintext mismatch (%d bytes back)", n, len(got))
		}
	}
}

func TestCiphertextDoesNotLeakPlaintext(t *testing.T) {
	secret := []byte("-----BEGIN RSA PRIVATE KEY-----\nMIIEow...\n")
	enc := encrypt(t, TypeDomainExport, "a very long password", secret)
	if bytes.Contains(enc, secret) {
		t.Fatal("plaintext found verbatim in the envelope")
	}
	if !HasMagic(enc) {
		t.Fatal("envelope does not start with the magic")
	}
}

func TestWrongPassword(t *testing.T) {
	enc := encrypt(t, TypeDomainExport, "the right password", []byte("secret payload"))
	_, _, err := decrypt(t, enc, "the wrong password")
	if !errors.Is(err, ErrWrongPassword) {
		t.Fatalf("err = %v, want ErrWrongPassword", err)
	}
}

func TestNotEncrypted(t *testing.T) {
	for _, plain := range [][]byte{[]byte(`{"format":"selfpost-domain-export"}`), {}, []byte("SELF")} {
		if _, err := NewReader(bytes.NewReader(plain), "pw"); !errors.Is(err, ErrNotEncrypted) {
			t.Errorf("NewReader(%q) = %v, want ErrNotEncrypted", plain, err)
		}
		if HasMagic(plain) {
			t.Errorf("HasMagic(%q) = true", plain)
		}
	}
}

func TestTruncationDetected(t *testing.T) {
	// A backup cut short mid-transfer must fail loudly rather than restore a
	// plausible-looking prefix.
	data := bytes.Repeat([]byte("payload"), 20000) // spans several chunks
	enc := encrypt(t, TypeFullBackup, "password password", data)

	// Drop the final chunk entirely: what remains is a sequence of valid,
	// correctly authenticated chunks, none of which is marked last.
	var chunks []int
	for off := headerLen; off < len(enc); {
		n := int(enc[off])<<24 | int(enc[off+1])<<16 | int(enc[off+2])<<8 | int(enc[off+3])
		chunks = append(chunks, off)
		off += 4 + n
	}
	if len(chunks) < 2 {
		t.Fatalf("expected several chunks, got %d", len(chunks))
	}
	truncated := enc[:chunks[len(chunks)-1]]
	r, err := NewReader(bytes.NewReader(truncated), "password password")
	if err != nil {
		t.Fatalf("NewReader (truncated): %v", err)
	}
	if _, err := io.ReadAll(r); !errors.Is(err, ErrCorrupt) {
		t.Fatalf("read truncated = %v, want ErrCorrupt", err)
	}
}

func TestTamperDetected(t *testing.T) {
	enc := encrypt(t, TypeFullBackup, "password password", []byte("some archive bytes"))
	for _, off := range []int{MagicLen /* type byte */, headerLen + 6 /* ciphertext */} {
		bad := bytes.Clone(enc)
		bad[off] ^= 0xff
		if _, _, err := decrypt(t, bad, "password password"); err == nil {
			t.Errorf("flipping byte %d was accepted", off)
		}
	}
}

func TestChunkReorderRejected(t *testing.T) {
	// Two full chunks plus a tail, so swapping the first two is possible without
	// changing any length.
	data := make([]byte, 2*chunkSize+10)
	for i := range data {
		data[i] = byte(i)
	}
	enc := encrypt(t, TypeFullBackup, "password password", data)

	const framed = 4 + chunkSize + tagLen
	first := headerLen
	second := headerLen + framed
	swapped := bytes.Clone(enc)
	copy(swapped[first:first+framed], enc[second:second+framed])
	copy(swapped[second:second+framed], enc[first:first+framed])

	if _, _, err := decrypt(t, swapped, "password password"); !errors.Is(err, ErrWrongPassword) {
		t.Fatalf("reordered chunks = %v, want authentication failure", err)
	}
}

func TestEmptyPasswordRejected(t *testing.T) {
	if _, err := NewWriter(io.Discard, TypeFullBackup, ""); err == nil {
		t.Fatal("empty password accepted")
	}
}

func TestUnreasonableKDFParamsRejected(t *testing.T) {
	enc := encrypt(t, TypeFullBackup, "password password", []byte("x"))
	bad := bytes.Clone(enc)
	bad[MagicLen+2] = 40 // logN far beyond the budget
	if _, err := NewReader(bytes.NewReader(bad), "password password"); !errors.Is(err, ErrCorrupt) {
		t.Fatalf("err = %v, want ErrCorrupt", err)
	}
}

// A header may ask for no more than the stated budget: 128·N·r bytes of
// scrypt's array at most 128 MiB, r at most 64, p at most 2 (audit
// 2026-10-12, M4; security review of stage 3). The writer's own parameters
// are a quarter of it.
func TestKDFBudget(t *testing.T) {
	enc := encrypt(t, TypeFullBackup, "password password", []byte("x"))
	header := func(logN uint8, r, p uint32) []byte {
		h := bytes.Clone(enc[:headerLen])
		h[MagicLen+2] = logN
		binary.BigEndian.PutUint32(h[MagicLen+3:], r)
		binary.BigEndian.PutUint32(h[MagicLen+7:], p)
		return h
	}
	for _, c := range []struct {
		logN uint8
		r, p uint32
		ok   bool
	}{
		{defaultLogN, defaultR, defaultP, true},
		{17, 8, 2, true},        // 128 MiB, the budget
		{20, 1, 1, true},        // 128 MiB another way
		{18, 8, 1, false},       // 256 MiB
		{15, 64, 1, false},      // 256 MiB
		{20, 32, 16, false},     // what the old cap let through: 4 GiB
		{15, 8, 3, false},       // p over 2
		{0, 8, 1, false},        // N of 1
		{31, 1, 1, false},       // past the shift
		{1, 1 << 21, 1, false},  // r alone over the budget
		{1, 1 << 19, 2, false},  // V small, B and XY 384 MiB
		{11, 64, 2, true},       // r at its cap
		{11, 65, 1, false},      // r past it
		{30, 1 << 20, 1, false}, // the largest product, still refused
		{15, 0, 1, false},
		{15, 8, 0, false},
	} {
		_, err := parseHeader(header(c.logN, c.r, c.p))
		if c.ok && err != nil {
			t.Errorf("logN=%d r=%d p=%d refused: %v", c.logN, c.r, c.p, err)
		}
		if !c.ok && !errors.Is(err, ErrCorrupt) {
			t.Errorf("logN=%d r=%d p=%d = %v, want ErrCorrupt", c.logN, c.r, c.p, err)
		}
	}
}

// Nothing may follow the last chunk: two envelopes end to end, or stray
// bytes, are damage rather than the first file (audit 2026-10-12, L3).
func TestBytesAfterTheLastChunkRejected(t *testing.T) {
	first := encrypt(t, TypeFullBackup, "password password", []byte("first"))
	second := encrypt(t, TypeFullBackup, "password password", []byte("second"))
	for name, data := range map[string][]byte{
		"two envelopes": append(bytes.Clone(first), second...),
		"one byte":      append(bytes.Clone(first), 0),
	} {
		if _, _, err := decrypt(t, data, "password password"); !errors.Is(err, ErrCorrupt) {
			t.Errorf("%s: err = %v, want ErrCorrupt", name, err)
		}
	}
	if plain, _, err := decrypt(t, first, "password password"); err != nil || string(plain) != "first" {
		t.Errorf("the envelope alone = %q, %v", plain, err)
	}
}