mailgoserver/internal/mailstore/mailstore_test.go

package mailstore

import (
	"bytes"
	"os"
	"path/filepath"
	"testing"

	"mailgoserver/internal/db"
)

func newTestDB(t *testing.T) *db.DB {
	t.Helper()
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "test.db")
	database, err := db.Open(dbPath)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { database.Close() })
	return database
}

// newTestMailbox creates a domain + mailbox with a real wrapped DEK, using a Store
// built against a random master key, and returns both.
func newTestMailbox(t *testing.T, quotaBytes int64) (*Store, int64) {
	t.Helper()
	database := newTestDB(t)
	domainID, err := database.CreateDomain("example.com")
	if err != nil {
		t.Fatal(err)
	}

	masterKey := GenerateDEK() // 32 random bytes, reused here as a throwaway master key
	s := New(database, masterKey, t.TempDir())

	dek := GenerateDEK()
	wrapped, nonce, err := s.WrapDEK(dek)
	if err != nil {
		t.Fatal(err)
	}
	hash, err := db.HashPassword("irrelevant-portal-password")
	if err != nil {
		t.Fatal(err)
	}
	mailboxID, err := database.CreateMailbox("user@example.com", hash, domainID, quotaBytes, wrapped, nonce)
	if err != nil {
		t.Fatal(err)
	}
	return s, mailboxID
}

func TestWrapUnwrapDEK(t *testing.T) {
	database := newTestDB(t)
	s := New(database, GenerateDEK(), t.TempDir())

	dek := GenerateDEK()
	wrapped, nonce, err := s.WrapDEK(dek)
	if err != nil {
		t.Fatal(err)
	}
	got, err := s.UnwrapDEK(wrapped, nonce)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(dek, got) {
		t.Fatalf("unwrapped DEK does not match original: got %x, want %x", got, dek)
	}
}

func TestStoreFetchRoundTrip(t *testing.T) {
	s, mailboxID := newTestMailbox(t, 1024*1024)
	raw := []byte("From: a@example.com\r\nSubject: hi\r\n\r\nhello world")

	uid, err := s.StoreMessage(mailboxID, "INBOX", raw, "<abc@example.com>", "a@example.com", "hi")
	if err != nil {
		t.Fatal(err)
	}

	got, err := s.FetchMessage(mailboxID, uid)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(raw, got) {
		t.Fatalf("fetched message does not match stored: got %q, want %q", got, raw)
	}

	msg, err := s.DB.GetMessageByUID(mailboxID, uid)
	if err != nil {
		t.Fatal(err)
	}
	onDisk, err := os.ReadFile(msg.StoragePath)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(onDisk, raw) {
		t.Fatal("on-disk file matches plaintext — message was not actually encrypted")
	}

	mbox, err := s.DB.GetMailboxByID(mailboxID)
	if err != nil {
		t.Fatal(err)
	}
	if mbox.UsedBytes != int64(len(raw)) {
		t.Fatalf("used_bytes = %d, want %d", mbox.UsedBytes, len(raw))
	}
}

func TestQuotaExceeded(t *testing.T) {
	s, mailboxID := newTestMailbox(t, 10) // tiny quota
	raw := []byte("this message is definitely longer than ten bytes")

	_, err := s.StoreMessage(mailboxID, "INBOX", raw, "<abc@example.com>", "a@example.com", "hi")
	if err != ErrQuotaExceeded {
		t.Fatalf("err = %v, want ErrQuotaExceeded", err)
	}

	mbox, err := s.DB.GetMailboxByID(mailboxID)
	if err != nil {
		t.Fatal(err)
	}
	if mbox.UsedBytes != 0 {
		t.Fatalf("used_bytes = %d after a rejected store, want 0", mbox.UsedBytes)
	}

	entries, err := os.ReadDir(s.BasePath)
	if err == nil && len(entries) != 0 {
		t.Fatalf("expected no files written under %s after a rejected store, found %d entries", s.BasePath, len(entries))
	}
}

func TestDeleteMessageFreesQuota(t *testing.T) {
	s, mailboxID := newTestMailbox(t, 1024*1024)
	raw := []byte("From: a@example.com\r\nSubject: bye\r\n\r\ngoodbye")

	uid, err := s.StoreMessage(mailboxID, "INBOX", raw, "<def@example.com>", "a@example.com", "bye")
	if err != nil {
		t.Fatal(err)
	}
	msg, err := s.DB.GetMessageByUID(mailboxID, uid)
	if err != nil {
		t.Fatal(err)
	}
	storagePath := msg.StoragePath

	if err := s.DeleteMessage(mailboxID, uid); err != nil {
		t.Fatal(err)
	}

	mbox, err := s.DB.GetMailboxByID(mailboxID)
	if err != nil {
		t.Fatal(err)
	}
	if mbox.UsedBytes != 0 {
		t.Fatalf("used_bytes = %d after delete, want 0", mbox.UsedBytes)
	}
	if _, err := os.Stat(storagePath); !os.IsNotExist(err) {
		t.Fatalf("ciphertext file %s still exists after delete", storagePath)
	}
}