encryption_test.go

package encryption

import (
	"bytes"
	"errors"
	"strconv"
	"testing"
)

func newTestBox(t *testing.T) *Box {
	t.Helper()
	box, err := New(make([]byte, KeySize))
	if err != nil {
		t.Fatalf("New() error = %v", err)
	}
	return box
}

func TestGenerateKey(t *testing.T) {
	first := GenerateKey()
	second := GenerateKey()
	if len(first) != KeySize || len(second) != KeySize {
		t.Fatalf("GenerateKey() lengths = %d and %d, want %d", len(first), len(second), KeySize)
	}
	if bytes.Equal(first, second) {
		t.Fatal("GenerateKey() returned identical keys")
	}
	if _, err := New(first); err != nil {
		t.Fatalf("New(GenerateKey()) error = %v", err)
	}
}

func TestNew(t *testing.T) {
	for _, size := range []int{0, 1, 16, 24, 31, 33, 64} {
		t.Run(strconv.Itoa(size), func(t *testing.T) {
			box, err := New(make([]byte, size))
			if box != nil || !errors.Is(err, ErrInvalidKeySize) {
				t.Fatalf("New(%d-byte key) = (%v, %v), want (nil, ErrInvalidKeySize)", size, box, err)
			}
		})
	}

	if box, err := New(make([]byte, KeySize)); err != nil || box == nil {
		t.Fatalf("New(%d-byte key) = (%v, %v), want non-nil Box", KeySize, box, err)
	}
}

func TestBoxRoundTrip(t *testing.T) {
	box := newTestBox(t)
	additionalData := []byte("tenant=example")
	for _, size := range []int{0, 1, 15, 16, 17, 31, 32, 33, 32 * 1024, 32*1024 + 1, 128 * 1024} {
		t.Run(strconv.Itoa(size), func(t *testing.T) {
			plaintext := bytes.Repeat([]byte{0xa5}, size)
			envelope := box.Seal(plaintext, additionalData)
			opened, err := box.Open(envelope, additionalData)
			if err != nil {
				t.Fatalf("Open() error = %v", err)
			}
			if !bytes.Equal(opened, plaintext) {
				t.Fatalf("Open(Seal(%d bytes)) returned %d different bytes", size, len(opened))
			}
		})
	}
}

func TestBoxSealUsesRandomNonce(t *testing.T) {
	box := newTestBox(t)
	plaintext := []byte("same plaintext")
	first := box.Seal(plaintext, nil)
	second := box.Seal(plaintext, nil)
	if bytes.Equal(first, second) {
		t.Fatal("Seal() returned identical envelopes for repeated plaintext")
	}
}

func TestEnvelopeFormat(t *testing.T) {
	box := newTestBox(t)
	plaintext := []byte("payload")
	envelope := box.Seal(plaintext, nil)
	if !bytes.Equal(envelope[:headerSize], envelopeHeader[:]) {
		t.Fatalf("Seal() header = %x, want %x", envelope[:headerSize], envelopeHeader)
	}
	if got, want := len(envelope), headerSize+box.aead.Overhead()+len(plaintext); got != want {
		t.Fatalf("len(Seal()) = %d, want %d", got, want)
	}
}

func TestBoxOpenRejectsInvalidInput(t *testing.T) {
	box := newTestBox(t)
	additionalData := []byte("context")
	envelope := box.Seal([]byte("secret"), additionalData)

	tests := map[string]struct {
		envelope       []byte
		additionalData []byte
	}{
		"nil":                   {nil, additionalData},
		"empty":                 {[]byte{}, additionalData},
		"header-only":           {append([]byte(nil), envelopeHeader[:]...), additionalData},
		"truncated":             {append([]byte(nil), envelope[:len(envelope)-1]...), additionalData},
		"appended":              {append(append([]byte(nil), envelope...), 0), additionalData},
		"wrong-additional-data": {append([]byte(nil), envelope...), []byte("different")},
		"wrong-magic":           {modifiedEnvelope(envelope, 0), additionalData},
		"unknown-version":       {modifiedEnvelope(envelope, 4), additionalData},
		"unknown-algorithm":     {modifiedEnvelope(envelope, 5), additionalData},
	}

	otherBox, err := New(bytes.Repeat([]byte{1}, KeySize))
	if err != nil {
		t.Fatalf("New(other key) error = %v", err)
	}
	if plaintext, err := otherBox.Open(envelope, additionalData); plaintext != nil || !errors.Is(err, ErrInvalidEnvelope) {
		t.Fatalf("Open() with wrong key = (%x, %v), want (nil, ErrInvalidEnvelope)", plaintext, err)
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			plaintext, err := box.Open(test.envelope, test.additionalData)
			if plaintext != nil || !errors.Is(err, ErrInvalidEnvelope) {
				t.Fatalf("Open() = (%x, %v), want (nil, ErrInvalidEnvelope)", plaintext, err)
			}
		})
	}
}

func modifiedEnvelope(envelope []byte, index int) []byte {
	modified := append([]byte(nil), envelope...)
	modified[index] ^= 0x01
	return modified
}

func TestBoxOpenRejectsModification(t *testing.T) {
	box := newTestBox(t)
	additionalData := []byte("context")
	envelope := box.Seal([]byte("secret payload"), additionalData)

	for index := range envelope {
		modified := append([]byte(nil), envelope...)
		modified[index] ^= 0x01
		plaintext, err := box.Open(modified, additionalData)
		if plaintext != nil || !errors.Is(err, ErrInvalidEnvelope) {
			t.Fatalf("Open() after modifying byte %d = (%x, %v), want (nil, ErrInvalidEnvelope)", index, plaintext, err)
		}
	}
}

func TestJoinAdditionalData(t *testing.T) {
	header := []byte("header")
	additionalData := []byte("context")
	joined := joinAdditionalData(header, additionalData)
	if got, want := string(joined), "headercontext"; got != want {
		t.Fatalf("joinAdditionalData() = %q, want %q", got, want)
	}
	joined[0] = 'H'
	if string(header) != "header" || string(additionalData) != "context" {
		t.Fatal("joinAdditionalData() aliases its inputs")
	}
}