Kyle Isom
564e0a9c67
Configurable maximum concurrent connections per listener. When the limit is reached, new connections are closed immediately after accept. 0 means unlimited (default, preserving existing behavior). Config: Listener gains max_connections field, validated non-negative. DB: Migration 3 adds listeners.max_connections column. UpdateListenerMaxConns method for runtime changes via gRPC. CreateListener updated to persist max_connections on seed. Server: ListenerState/ListenerData gain MaxConnections. Limit checked in serve() after Accept but before handleConn — if ActiveConnections >= MaxConnections, connection is closed and the accept loop continues. SetMaxConnections method for runtime updates. Proto: SetListenerMaxConnections RPC added. ListenerStatus gains max_connections field. Generated code regenerated. gRPC server: SetListenerMaxConnections implements write-through (DB first, then in-memory update). GetStatus includes max_connections. Client: SetListenerMaxConnections method, MaxConnections in ListenerStatus. Tests: DB CRUD and UpdateListenerMaxConns, server connection limit enforcement (accept 2, reject 3rd, close one, accept again), gRPC SetListenerMaxConnections round-trip with DB persistence, not-found error handling. Also updates PROJECT_PLAN.md with phases 6-8 and PROGRESS.md with tracking for the new features. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
511 lines
12 KiB
Go
511 lines
12 KiB
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, 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 TestRouteDuplicateHostname(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)
|
|
}
|
|
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, 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 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")
|
|
}
|
|
}
|