Files
mc-proxy/internal/db/db_test.go
Kyle Isom 28321e22f4 Make AddRoute idempotent (upsert instead of reject duplicates)
AddRoute now updates an existing route if one already exists for the
same (listener, hostname) pair, instead of returning AlreadyExists.
This makes repeated deploys idempotent — the MCP agent can register
routes on every deploy without needing to remove them first.

- DB: INSERT ... ON CONFLICT DO UPDATE (SQLite upsert)
- In-memory: overwrite existing route unconditionally
- gRPC: error code changed from AlreadyExists to Internal (for real DB errors)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-27 14:01:45 -07:00

618 lines
15 KiB
Go

package db
import (
"path/filepath"
"testing"
"git.wntrmute.dev/mc/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, 0); 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, 0)
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, 0)
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 TestListenerMaxConnections(t *testing.T) {
store := openTestDB(t)
id, err := store.CreateListener(":443", false, 5000)
if err != nil {
t.Fatalf("create: %v", err)
}
l, err := store.GetListenerByAddr(":443")
if err != nil {
t.Fatalf("get: %v", err)
}
if l.MaxConnections != 5000 {
t.Fatalf("max_connections = %d, want 5000", l.MaxConnections)
}
// Update max connections.
if err := store.UpdateListenerMaxConns(id, 10000); err != nil {
t.Fatalf("update: %v", err)
}
l, err = store.GetListenerByAddr(":443")
if err != nil {
t.Fatalf("get after update: %v", err)
}
if l.MaxConnections != 10000 {
t.Fatalf("max_connections = %d, want 10000", l.MaxConnections)
}
// Set to 0 (unlimited).
if err := store.UpdateListenerMaxConns(id, 0); err != nil {
t.Fatalf("update to 0: %v", err)
}
l, _ = store.GetListenerByAddr(":443")
if l.MaxConnections != 0 {
t.Fatalf("max_connections = %d, want 0", l.MaxConnections)
}
}
func TestListenerDuplicateAddr(t *testing.T) {
store := openTestDB(t)
if _, err := store.CreateListener(":443", false, 0); err != nil {
t.Fatalf("first create: %v", err)
}
if _, err := store.CreateListener(":443", false, 0); err == nil {
t.Fatal("expected error for duplicate addr")
}
}
func TestRouteCRUD(t *testing.T) {
store := openTestDB(t)
listenerID, err := store.CreateListener(":443", false, 0)
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, 0)
_, 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 TestRouteUpsert(t *testing.T) {
store := openTestDB(t)
listenerID, _ := store.CreateListener(":443", false, 0)
if _, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:8443", "l4", "", "", false, false); err != nil {
t.Fatalf("first create: %v", err)
}
// Same (listener, hostname) with different backend — should upsert, not error.
if _, err := store.CreateRoute(listenerID, "example.com", "127.0.0.1:9443", "l7", "/cert.pem", "/key.pem", false, false); err != nil {
t.Fatalf("upsert: %v", err)
}
routes, err := store.ListRoutes(listenerID)
if err != nil {
t.Fatalf("list routes: %v", err)
}
if len(routes) != 1 {
t.Fatalf("expected 1 route after upsert, got %d", len(routes))
}
if routes[0].Backend != "127.0.0.1:9443" {
t.Fatalf("expected updated backend, got %q", routes[0].Backend)
}
if routes[0].Mode != "l7" {
t.Fatalf("expected updated mode, got %q", routes[0].Mode)
}
}
func TestRouteCascadeDelete(t *testing.T) {
store := openTestDB(t)
listenerID, _ := store.CreateListener(":443", false, 0)
_, _ = 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, 0)
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 func() { _ = 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, 0)
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")
}
}
func TestL7PolicyCRUD(t *testing.T) {
store := openTestDB(t)
lid, _ := store.CreateListener(":443", false, 0)
rid, _ := store.CreateRoute(lid, "api.test", "127.0.0.1:8080", "l7", "/c.pem", "/k.pem", false, false)
// Create policies.
id1, err := store.CreateL7Policy(rid, "block_user_agent", "BadBot")
if err != nil {
t.Fatalf("create policy 1: %v", err)
}
if id1 == 0 {
t.Fatal("expected non-zero policy ID")
}
if _, err := store.CreateL7Policy(rid, "require_header", "X-API-Key"); err != nil {
t.Fatalf("create policy 2: %v", err)
}
// List policies.
policies, err := store.ListL7Policies(rid)
if err != nil {
t.Fatalf("list: %v", err)
}
if len(policies) != 2 {
t.Fatalf("got %d policies, want 2", len(policies))
}
// Delete one.
if err := store.DeleteL7Policy(rid, "block_user_agent", "BadBot"); err != nil {
t.Fatalf("delete: %v", err)
}
policies, _ = store.ListL7Policies(rid)
if len(policies) != 1 {
t.Fatalf("got %d policies after delete, want 1", len(policies))
}
if policies[0].Type != "require_header" {
t.Fatalf("remaining policy type = %q, want %q", policies[0].Type, "require_header")
}
}
func TestL7PolicyCascadeDelete(t *testing.T) {
store := openTestDB(t)
lid, _ := store.CreateListener(":443", false, 0)
rid, _ := store.CreateRoute(lid, "api.test", "127.0.0.1:8080", "l7", "/c.pem", "/k.pem", false, false)
_, _ = store.CreateL7Policy(rid, "block_user_agent", "Bot")
// Deleting the route should cascade-delete its policies.
_ = store.DeleteRoute(lid, "api.test")
policies, _ := store.ListL7Policies(rid)
if len(policies) != 0 {
t.Fatalf("got %d policies after cascade delete, want 0", len(policies))
}
}
func TestL7PolicyDuplicate(t *testing.T) {
store := openTestDB(t)
lid, _ := store.CreateListener(":443", false, 0)
rid, _ := store.CreateRoute(lid, "api.test", "127.0.0.1:8080", "l7", "/c.pem", "/k.pem", false, false)
if _, err := store.CreateL7Policy(rid, "block_user_agent", "Bot"); err != nil {
t.Fatalf("first create: %v", err)
}
if _, err := store.CreateL7Policy(rid, "block_user_agent", "Bot"); err == nil {
t.Fatal("expected error for duplicate policy")
}
}
func TestGetRouteID(t *testing.T) {
store := openTestDB(t)
lid, _ := store.CreateListener(":443", false, 0)
_, _ = store.CreateRoute(lid, "api.test", "127.0.0.1:8080", "l7", "/c.pem", "/k.pem", false, false)
rid, err := store.GetRouteID(lid, "api.test")
if err != nil {
t.Fatalf("GetRouteID: %v", err)
}
if rid == 0 {
t.Fatal("expected non-zero route ID")
}
_, err = store.GetRouteID(lid, "nonexistent.test")
if err == nil {
t.Fatal("expected error for nonexistent route")
}
}