mc-proxy/internal/db/db_test.go

package db

import (
	"path/filepath"
	"testing"

	"git.wntrmute.dev/kyle/mc-proxy/internal/config"
)

func openTestDB(t *testing.T) *Store {
	t.Helper()
	dir := t.TempDir()
	store, err := Open(filepath.Join(dir, "test.db"))
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	if err := store.Migrate(); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	t.Cleanup(func() { store.Close() })
	return store
}

func TestMigrate(t *testing.T) {
	store := openTestDB(t)

	// Running migrate again should be idempotent.
	if err := store.Migrate(); err != nil {
		t.Fatalf("second migrate: %v", err)
	}
}

func TestIsEmpty(t *testing.T) {
	store := openTestDB(t)

	empty, err := store.IsEmpty()
	if err != nil {
		t.Fatalf("is empty: %v", err)
	}
	if !empty {
		t.Fatal("expected empty database")
	}

	if _, err := store.CreateListener(":443"); err != nil {
		t.Fatalf("create listener: %v", err)
	}

	empty, err = store.IsEmpty()
	if err != nil {
		t.Fatalf("is empty: %v", err)
	}
	if empty {
		t.Fatal("expected non-empty database")
	}
}

func TestListenerCRUD(t *testing.T) {
	store := openTestDB(t)

	id, err := store.CreateListener(":443")
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	if id == 0 {
		t.Fatal("expected non-zero ID")
	}

	listeners, err := store.ListListeners()
	if err != nil {
		t.Fatalf("list: %v", err)
	}
	if len(listeners) != 1 {
		t.Fatalf("got %d listeners, want 1", len(listeners))
	}
	if listeners[0].Addr != ":443" {
		t.Fatalf("got addr %q, want %q", listeners[0].Addr, ":443")
	}

	l, err := store.GetListenerByAddr(":443")
	if err != nil {
		t.Fatalf("get by addr: %v", err)
	}
	if l.ID != id {
		t.Fatalf("got ID %d, want %d", l.ID, id)
	}

	if err := store.DeleteListener(id); err != nil {
		t.Fatalf("delete: %v", err)
	}

	listeners, err = store.ListListeners()
	if err != nil {
		t.Fatalf("list after delete: %v", err)
	}
	if len(listeners) != 0 {
		t.Fatalf("got %d listeners after delete, want 0", len(listeners))
	}
}

func TestListenerDuplicateAddr(t *testing.T) {
	store := openTestDB(t)

	if _, err := store.CreateListener(":443"); err != nil {
		t.Fatalf("first create: %v", err)
	}
	if _, err := store.CreateListener(":443"); err == nil {
		t.Fatal("expected error for duplicate addr")
	}
}

func TestRouteCRUD(t *testing.T) {
	store := openTestDB(t)

	listenerID, err := store.CreateListener(":443")
	if err != nil {
		t.Fatalf("create listener: %v", err)
	}

	routeID, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:8443")
	if err != nil {
		t.Fatalf("create route: %v", err)
	}
	if routeID == 0 {
		t.Fatal("expected non-zero route ID")
	}

	routes, err := store.ListRoutes(listenerID)
	if err != nil {
		t.Fatalf("list routes: %v", err)
	}
	if len(routes) != 1 {
		t.Fatalf("got %d routes, want 1", len(routes))
	}
	if routes[0].Hostname != "example.com" {
		t.Fatalf("got hostname %q, want %q", routes[0].Hostname, "example.com")
	}

	if err := store.DeleteRoute(listenerID, "example.com"); err != nil {
		t.Fatalf("delete route: %v", err)
	}

	routes, err = store.ListRoutes(listenerID)
	if err != nil {
		t.Fatalf("list after delete: %v", err)
	}
	if len(routes) != 0 {
		t.Fatalf("got %d routes after delete, want 0", len(routes))
	}
}

func TestRouteDuplicateHostname(t *testing.T) {
	store := openTestDB(t)

	listenerID, _ := store.CreateListener(":443")
	if _, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:8443"); err != nil {
		t.Fatalf("first create: %v", err)
	}
	if _, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:9443"); err == nil {
		t.Fatal("expected error for duplicate hostname on same listener")
	}
}

func TestRouteCascadeDelete(t *testing.T) {
	store := openTestDB(t)

	listenerID, _ := store.CreateListener(":443")
	store.CreateRoute(listenerID, "a.example.com", "127.0.0.1:8443")
	store.CreateRoute(listenerID, "b.example.com", "127.0.0.1:9443")

	if err := store.DeleteListener(listenerID); err != nil {
		t.Fatalf("delete listener: %v", err)
	}

	routes, err := store.ListRoutes(listenerID)
	if err != nil {
		t.Fatalf("list routes: %v", err)
	}
	if len(routes) != 0 {
		t.Fatalf("got %d routes after cascade delete, want 0", len(routes))
	}
}

func TestFirewallRuleCRUD(t *testing.T) {
	store := openTestDB(t)

	id, err := store.CreateFirewallRule("ip", "192.0.2.1")
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	if id == 0 {
		t.Fatal("expected non-zero ID")
	}

	if _, err := store.CreateFirewallRule("cidr", "198.51.100.0/24"); err != nil {
		t.Fatalf("create cidr: %v", err)
	}
	if _, err := store.CreateFirewallRule("country", "CN"); err != nil {
		t.Fatalf("create country: %v", err)
	}

	rules, err := store.ListFirewallRules()
	if err != nil {
		t.Fatalf("list: %v", err)
	}
	if len(rules) != 3 {
		t.Fatalf("got %d rules, want 3", len(rules))
	}

	if err := store.DeleteFirewallRule("ip", "192.0.2.1"); err != nil {
		t.Fatalf("delete: %v", err)
	}

	rules, err = store.ListFirewallRules()
	if err != nil {
		t.Fatalf("list after delete: %v", err)
	}
	if len(rules) != 2 {
		t.Fatalf("got %d rules after delete, want 2", len(rules))
	}
}

func TestFirewallRuleDuplicate(t *testing.T) {
	store := openTestDB(t)

	if _, err := store.CreateFirewallRule("ip", "192.0.2.1"); err != nil {
		t.Fatalf("first create: %v", err)
	}
	if _, err := store.CreateFirewallRule("ip", "192.0.2.1"); err == nil {
		t.Fatal("expected error for duplicate rule")
	}
}

func TestSeed(t *testing.T) {
	store := openTestDB(t)

	listeners := []config.Listener{
		{
			Addr: ":443",
			Routes: []config.Route{
				{Hostname: "a.example.com", Backend: "127.0.0.1:8443"},
				{Hostname: "b.example.com", Backend: "127.0.0.1:9443"},
			},
		},
		{
			Addr: ":8443",
			Routes: []config.Route{
				{Hostname: "c.example.com", Backend: "127.0.0.1:18443"},
			},
		},
	}

	fw := config.Firewall{
		BlockedIPs:       []string{"192.0.2.1"},
		BlockedCIDRs:     []string{"198.51.100.0/24"},
		BlockedCountries: []string{"cn", "KP"},
	}

	if err := store.Seed(listeners, fw); err != nil {
		t.Fatalf("seed: %v", err)
	}

	dbListeners, err := store.ListListeners()
	if err != nil {
		t.Fatalf("list listeners: %v", err)
	}
	if len(dbListeners) != 2 {
		t.Fatalf("got %d listeners, want 2", len(dbListeners))
	}

	routes, err := store.ListRoutes(dbListeners[0].ID)
	if err != nil {
		t.Fatalf("list routes: %v", err)
	}
	if len(routes) != 2 {
		t.Fatalf("got %d routes for listener 0, want 2", len(routes))
	}

	rules, err := store.ListFirewallRules()
	if err != nil {
		t.Fatalf("list firewall rules: %v", err)
	}
	if len(rules) != 4 {
		t.Fatalf("got %d firewall rules, want 4", len(rules))
	}
}

func TestSnapshot(t *testing.T) {
	store := openTestDB(t)

	store.CreateListener(":443")

	dest := filepath.Join(t.TempDir(), "backup.db")
	if err := store.Snapshot(dest); err != nil {
		t.Fatalf("snapshot: %v", err)
	}

	// Open the snapshot and verify.
	backup, err := Open(dest)
	if err != nil {
		t.Fatalf("open backup: %v", err)
	}
	defer backup.Close()

	if err := backup.Migrate(); err != nil {
		t.Fatalf("migrate backup: %v", err)
	}

	listeners, err := backup.ListListeners()
	if err != nil {
		t.Fatalf("list from backup: %v", err)
	}
	if len(listeners) != 1 {
		t.Fatalf("backup has %d listeners, want 1", len(listeners))
	}
}

func TestDeleteNonexistent(t *testing.T) {
	store := openTestDB(t)

	if err := store.DeleteListener(999); err == nil {
		t.Fatal("expected error deleting nonexistent listener")
	}

	if err := store.DeleteRoute(999, "example.com"); err == nil {
		t.Fatal("expected error deleting nonexistent route")
	}

	if err := store.DeleteFirewallRule("ip", "1.2.3.4"); err == nil {
		t.Fatal("expected error deleting nonexistent firewall rule")
	}
}