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", false); 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", false)
	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")
	}
	if listeners[0].ProxyProtocol {
		t.Fatal("expected proxy_protocol = false")
	}

	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 TestListenerProxyProtocol(t *testing.T) {
	store := openTestDB(t)

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

	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 !l.ProxyProtocol {
		t.Fatal("expected proxy_protocol = true")
	}
}

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

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

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

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

	routeID, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:8443", "l4", "", "", false, false)
	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 routes[0].Mode != "l4" {
		t.Fatalf("got mode %q, want %q", routes[0].Mode, "l4")
	}

	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 TestRouteL7Fields(t *testing.T) {
	store := openTestDB(t)

	listenerID, _ := store.CreateListener(":443", false)

	_, err := store.CreateRoute(listenerID, "api.example.com", "127.0.0.1:8080", "l7",
		"/certs/api.crt", "/certs/api.key", false, true)
	if err != nil {
		t.Fatalf("create L7 route: %v", err)
	}

	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))
	}

	r := routes[0]
	if r.Mode != "l7" {
		t.Fatalf("mode = %q, want %q", r.Mode, "l7")
	}
	if r.TLSCert != "/certs/api.crt" {
		t.Fatalf("tls_cert = %q, want %q", r.TLSCert, "/certs/api.crt")
	}
	if r.TLSKey != "/certs/api.key" {
		t.Fatalf("tls_key = %q, want %q", r.TLSKey, "/certs/api.key")
	}
	if r.BackendTLS {
		t.Fatal("expected backend_tls = false")
	}
	if !r.SendProxyProtocol {
		t.Fatal("expected send_proxy_protocol = true")
	}
}

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

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

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

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

	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", Mode: "l4"},
				{Hostname: "b.example.com", Backend: "127.0.0.1:9443"},
			},
		},
		{
			Addr:          ":8443",
			ProxyProtocol: true,
			Routes: []config.Route{
				{Hostname: "c.example.com", Backend: "127.0.0.1:18443", Mode: "l4", SendProxyProtocol: true},
			},
		},
	}

	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))
	}
	if !dbListeners[1].ProxyProtocol {
		t.Fatal("expected listener 2 proxy_protocol = true")
	}

	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))
	}

	// Verify mode defaults to "l4" even when empty in config.
	for _, r := range routes {
		if r.Mode != "l4" {
			t.Fatalf("route %q mode = %q, want %q", r.Hostname, r.Mode, "l4")
		}
	}

	// Verify send_proxy_protocol on listener 2's route.
	routes2, err := store.ListRoutes(dbListeners[1].ID)
	if err != nil {
		t.Fatalf("list routes listener 2: %v", err)
	}
	if len(routes2) != 1 {
		t.Fatalf("got %d routes for listener 1, want 1", len(routes2))
	}
	if !routes2[0].SendProxyProtocol {
		t.Fatal("expected send_proxy_protocol = true on listener 2 route")
	}

	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", false)

	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")
	}
}

// TestMigrationV2Upgrade verifies that migration v2 adds new columns
// to an existing v1 database without data loss.
func TestMigrationV2Upgrade(t *testing.T) {
	dir := t.TempDir()
	store, err := Open(filepath.Join(dir, "test.db"))
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	t.Cleanup(func() { store.Close() })

	// Run full migrations (v1 + v2).
	if err := store.Migrate(); err != nil {
		t.Fatalf("migrate: %v", err)
	}

	// Insert a listener and route with defaults to verify new columns work.
	lid, err := store.CreateListener(":443", false)
	if err != nil {
		t.Fatalf("create listener: %v", err)
	}

	_, err = store.CreateRoute(lid, "test.example.com", "127.0.0.1:8443", "l4", "", "", false, false)
	if err != nil {
		t.Fatalf("create route: %v", err)
	}

	// Read back and verify defaults.
	routes, err := store.ListRoutes(lid)
	if err != nil {
		t.Fatalf("list routes: %v", err)
	}
	if len(routes) != 1 {
		t.Fatalf("got %d routes, want 1", len(routes))
	}
	r := routes[0]
	if r.Mode != "l4" {
		t.Fatalf("mode = %q, want %q", r.Mode, "l4")
	}
	if r.TLSCert != "" || r.TLSKey != "" {
		t.Fatalf("expected empty cert/key, got cert=%q key=%q", r.TLSCert, r.TLSKey)
	}
	if r.BackendTLS || r.SendProxyProtocol {
		t.Fatal("expected false for backend_tls and send_proxy_protocol")
	}

	listeners, err := store.ListListeners()
	if err != nil {
		t.Fatalf("list listeners: %v", err)
	}
	if listeners[0].ProxyProtocol {
		t.Fatal("expected proxy_protocol = false")
	}
}