diff --git a/server/cmd/server/main.go b/server/cmd/server/main.go index 37d7965..b78c764 100644 --- a/server/cmd/server/main.go +++ b/server/cmd/server/main.go @@ -1,20 +1,31 @@ package main import ( + "context" + "crypto/tls" "encoding/json" "flag" "log" + "log/slog" + "net" "net/http" "os" "github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5/middleware" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials" + + "github.com/wangjia/pangolin/server/internal/db" + "github.com/wangjia/pangolin/server/internal/mtls" + "github.com/wangjia/pangolin/server/internal/nodes" + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" "github.com/wangjia/pangolin/server/internal/redisutil" "github.com/wangjia/pangolin/server/internal/scheduler/probe" ) func main() { - addr := flag.String("addr", "", "listen address (default :8080, overridden by ADDR env)") + addr := flag.String("addr", "", "HTTP listen address (default :8080, overridden by ADDR env)") flag.Parse() if *addr == "" { @@ -25,6 +36,8 @@ func main() { } } + // ─── HTTP server ────────────────────────────────────────────────────────── + r := chi.NewRouter() r.Use(middleware.Logger) r.Use(middleware.Recoverer) @@ -36,14 +49,6 @@ func main() { }) // --- Probe ingest route (optional) --- - // Requires: - // PROBE_SECRETS – JSON object mapping probeId → HMAC secret, e.g. - // '{"probe-sg-01":"s3cr3t1","probe-jp-01":"s3cr3t2"}' - // REDIS_ADDR – Redis address, default 127.0.0.1:6379 - // REDIS_PASSWORD – Redis password (optional) - // - // If PROBE_SECRETS is empty the route is not registered and the server - // starts normally without the probe ingest endpoint. probeSecretsJSON := os.Getenv("PROBE_SECRETS") if probeSecretsJSON != "" { redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379") @@ -63,12 +68,101 @@ func main() { } } + // ─── gRPC agent server (optional) ──────────────────────────────────────── + // All five env vars must be set to activate the gRPC listener. + + grpcAddr := os.Getenv("GRPC_ADDR") + caKeyPath := os.Getenv("CA_KEY_PATH") + caCertPath := os.Getenv("CA_CERT_PATH") + grpcCertPath := os.Getenv("GRPC_CERT_PATH") + grpcKeyPath := os.Getenv("GRPC_KEY_PATH") + dbDSN := os.Getenv("DB_DSN") + + if grpcAddr != "" && caKeyPath != "" && caCertPath != "" && + grpcCertPath != "" && grpcKeyPath != "" && dbDSN != "" { + redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379") + go startGRPC(grpcAddr, caKeyPath, caCertPath, grpcCertPath, grpcKeyPath, + redisAddr, os.Getenv("REDIS_PASSWORD"), dbDSN) + } else if grpcAddr != "" { + slog.Warn("grpc server disabled: one or more required env vars are missing", + "GRPC_ADDR", grpcAddr, + "CA_KEY_PATH_set", caKeyPath != "", + "CA_CERT_PATH_set", caCertPath != "", + "GRPC_CERT_PATH_set", grpcCertPath != "", + "GRPC_KEY_PATH_set", grpcKeyPath != "", + "DB_DSN_set", dbDSN != "") + } + + // ─── Start HTTP server (blocking) ──────────────────────────────────────── + log.Printf("pangolin server listening on %s", *addr) if err := http.ListenAndServe(*addr, r); err != nil { log.Fatalf("server error: %v", err) } } +// startGRPC initialises and starts the mTLS gRPC agent server. +// Must be called in a goroutine. Calls log.Fatalf on unrecoverable errors. +func startGRPC( + addr, caKeyPath, caCertPath, grpcCertPath, grpcKeyPath, + redisAddr, redisPassword, dbDSN string, +) { + // Pangolin Node CA (signs + verifies node client certs). + ca, err := mtls.NewCA(mtls.CAConfig{ + KeyPath: caKeyPath, + CertPath: caCertPath, + }) + if err != nil { + log.Fatalf("grpc: load CA: %v", err) + } + + // Redis (hub command queue, load cache, CRL, bootstrap tokens). + rdb, err := redisutil.New(redisAddr, redisPassword, 0) + if err != nil { + log.Fatalf("grpc: redis connect: %v", err) + } + + // CRL + one-time bootstrap token manager. + crl := mtls.NewCRL(rdb, nil) + tokens := mtls.NewBootstrapTokenManager(rdb) + + // MySQL for node catalogue and usage accumulation. + sqlDB, err := db.Open(dbDSN) + if err != nil { + log.Fatalf("grpc: db open: %v", err) + } + + // Nodes service (wires Hub, LoadCache, Handler). + store := nodes.NewSQLNodeStore(sqlDB) + svc := nodes.NewService(ca, tokens, rdb, store) + svc.Hub().Start(context.Background()) // pub/sub goroutine runs for process lifetime + + // gRPC transport TLS (server's own cert, typically from Let's Encrypt). + serverCert, err := tls.LoadX509KeyPair(grpcCertPath, grpcKeyPath) + if err != nil { + log.Fatalf("grpc: load transport TLS cert: %v", err) + } + tlsCfg := mtls.NewServerTLSConfig(ca, crl) + tlsCfg.Certificates = []tls.Certificate{serverCert} + + // Build gRPC server with mTLS transport + identity interceptors. + srv := grpc.NewServer( + grpc.Creds(credentials.NewTLS(tlsCfg)), + grpc.ChainUnaryInterceptor(mtls.UnaryServerInterceptor()), + grpc.ChainStreamInterceptor(mtls.StreamServerInterceptor()), + ) + agentv1.RegisterAgentServiceServer(srv, svc.Handler()) + + lis, err := net.Listen("tcp", addr) + if err != nil { + log.Fatalf("grpc: listen %s: %v", addr, err) + } + slog.Info("grpc agent server listening", "addr", addr) + if err := srv.Serve(lis); err != nil { + log.Fatalf("grpc: serve: %v", err) + } +} + func getenvDefault(key, def string) string { if v := os.Getenv(key); v != "" { return v diff --git a/server/internal/config/config.go b/server/internal/config/config.go index fcfd5e5..7345445 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -37,6 +37,24 @@ type Config struct { // WebhookNonceTTL is how long a webhook nonce is kept in Redis to prevent replay. // Should be > 2 * WebhookTimestampTolerance. Default: 15 minutes. WebhookNonceTTL time.Duration + + // ─── gRPC / agent server ────────────────────────────────────────────── + + // GRPCAddr is the listen address for the mTLS gRPC agent server (host:port). + // Optional: if empty, the gRPC server is not started. + // Example: ":9443" + GRPCAddr string + + // CAKeyPath and CACertPath are file paths for the Pangolin Node CA. + // Required when GRPCAddr is set. + CAKeyPath string + CACertPath string + + // GRPCCertPath and GRPCKeyPath are the TLS certificate/key used by the gRPC + // listener for its own transport credential (distinct from the node CA). + // Required when GRPCAddr is set. + GRPCCertPath string + GRPCKeyPath string } // FromEnv reads configuration from environment variables. @@ -52,6 +70,12 @@ func FromEnv() (*Config, error) { RedeemLockDuration: time.Hour, WebhookTimestampTolerance: 5 * time.Minute, WebhookNonceTTL: 15 * time.Minute, + // gRPC server (optional) + GRPCAddr: os.Getenv("GRPC_ADDR"), + CAKeyPath: os.Getenv("CA_KEY_PATH"), + CACertPath: os.Getenv("CA_CERT_PATH"), + GRPCCertPath: os.Getenv("GRPC_CERT_PATH"), + GRPCKeyPath: os.Getenv("GRPC_KEY_PATH"), } if c.DSN == "" { diff --git a/server/internal/nodes/grpc_test.go b/server/internal/nodes/grpc_test.go new file mode 100644 index 0000000..755cc9e --- /dev/null +++ b/server/internal/nodes/grpc_test.go @@ -0,0 +1,819 @@ +package nodes_test + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "net" + "sync" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + + "github.com/wangjia/pangolin/server/internal/mtls" + "github.com/wangjia/pangolin/server/internal/nodes" + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +// ─── mock NodeStore ─────────────────────────────────────────────────────────── + +// mockNodeStore is an in-memory NodeStore used by gRPC integration tests. +type mockNodeStore struct { + mu sync.Mutex + nodeUUID string + configVer int64 + usageAccum []mockUsageEntry + usersByDpUUID map[string]int64 +} + +type mockUsageEntry struct { + UserID int64 + Date time.Time + BytesUp int64 + BytesDown int64 + Minutes int64 +} + +func (m *mockNodeStore) NodeByUUID(_ context.Context, uuid string) (*nodes.NodeRow, error) { + if uuid == m.nodeUUID { + return &nodes.NodeRow{ + ID: 1, + UUID: uuid, + Status: "up", + RealityPBK: "test-private-key", + RealitySNI: "www.example.com", + Endpoint: "1.2.3.4:443", + }, nil + } + return nil, nil +} + +func (m *mockNodeStore) ConfigVersion(_ context.Context) (int64, error) { + return m.configVer, nil +} + +func (m *mockNodeStore) ActiveNodeUUIDs(_ context.Context) ([]string, error) { + return []string{m.nodeUUID}, nil +} + +func (m *mockNodeStore) CredentialsForNode(_ context.Context, _ string) ([]*agentv1.Credential, error) { + return nil, nil +} + +func (m *mockNodeStore) UserIDByDpUUID(_ context.Context, dpUUID string) (int64, bool, error) { + if uid, ok := m.usersByDpUUID[dpUUID]; ok { + return uid, true, nil + } + return 0, false, nil +} + +func (m *mockNodeStore) AccumulateUsage(_ context.Context, userID int64, date time.Time, + bytesUp, bytesDown, minutes int64, +) error { + m.mu.Lock() + defer m.mu.Unlock() + m.usageAccum = append(m.usageAccum, mockUsageEntry{userID, date, bytesUp, bytesDown, minutes}) + return nil +} + +func (m *mockNodeStore) usageLog() []mockUsageEntry { + m.mu.Lock() + defer m.mu.Unlock() + out := make([]mockUsageEntry, len(m.usageAccum)) + copy(out, m.usageAccum) + return out +} + +// ─── test infrastructure helpers ───────────────────────────────────────────── + +// makeTestServerCert generates a self-signed ECDSA P-256 TLS certificate +// for use as the gRPC server's transport credential. +func makeTestServerCert(t *testing.T) tls.Certificate { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("gen server key: %v", err) + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "pangolin-grpc-test"}, + NotBefore: time.Now().Add(-time.Minute), + NotAfter: time.Now().Add(time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + BasicConstraintsValid: true, + } + certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + t.Fatalf("create server cert: %v", err) + } + keyDER, err := x509.MarshalECPrivateKey(key) + if err != nil { + t.Fatalf("marshal server key: %v", err) + } + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER}) + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}) + cert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + t.Fatalf("x509 keypair server: %v", err) + } + return cert +} + +// genKeyAndCSR generates an ECDSA P-256 private key and a PKCS#10 CSR (both PEM). +func genKeyAndCSR(t *testing.T) (keyPEM, csrPEM []byte) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("gen node key: %v", err) + } + keyDER, err := x509.MarshalECPrivateKey(key) + if err != nil { + t.Fatalf("marshal node key: %v", err) + } + keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}) + + tmpl := &x509.CertificateRequest{ + Subject: pkix.Name{CommonName: "pangolin-node-pending"}, + SignatureAlgorithm: x509.ECDSAWithSHA256, + } + csrDER, err := x509.CreateCertificateRequest(rand.Reader, tmpl, key) + if err != nil { + t.Fatalf("create CSR: %v", err) + } + csrPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrDER}) + return +} + +// testServerBundle holds all components for a single test gRPC server instance. +type testServerBundle struct { + t *testing.T + lis *bufconn.Listener + srv *grpc.Server + hub *nodes.Hub + rdb *redis.Client + store *mockNodeStore + tokens *mtls.BootstrapTokenManager + load *nodes.LoadCache +} + +// newTestServer creates a fully wired gRPC server (bufconn, miniredis, mTLS). +func newTestServer(t *testing.T, configVer int64, nodeUUID string) *testServerBundle { + t.Helper() + + // miniredis + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { rdb.Close() }) + + // CA + CRL + token manager + dir := t.TempDir() + ca, err := mtls.NewCA(mtls.CAConfig{ + KeyPath: dir + "/ca.key", + CertPath: dir + "/ca.crt", + }) + if err != nil { + t.Fatalf("NewCA: %v", err) + } + crl := mtls.NewCRL(rdb, nil) + tokens := mtls.NewBootstrapTokenManager(rdb) + + // Mock store + store := &mockNodeStore{ + nodeUUID: nodeUUID, + configVer: configVer, + usersByDpUUID: map[string]int64{"dp-user1": 101}, + } + + // Hub + load cache + hub := nodes.NewHub(rdb) + load := nodes.NewLoadCache(rdb) + ctx, cancel := context.WithCancel(context.Background()) + hub.Start(ctx) + t.Cleanup(cancel) + + // Handler + handler := nodes.NewHandler(ca, tokens, hub, store, load) + + // gRPC server TLS config + serverCert := makeTestServerCert(t) + tlsCfg := mtls.NewServerTLSConfig(ca, crl) + tlsCfg.Certificates = []tls.Certificate{serverCert} + + // gRPC server + lis := bufconn.Listen(1 << 20) + srv := grpc.NewServer( + grpc.Creds(credentials.NewTLS(tlsCfg)), + grpc.ChainUnaryInterceptor(mtls.UnaryServerInterceptor()), + grpc.ChainStreamInterceptor(mtls.StreamServerInterceptor()), + ) + agentv1.RegisterAgentServiceServer(srv, handler) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + + return &testServerBundle{ + t: t, lis: lis, srv: srv, hub: hub, + rdb: rdb, store: store, tokens: tokens, load: load, + } +} + +// bufDialer returns a context dialer that routes through the bufconn listener. +func bufDialer(lis *bufconn.Listener) func(context.Context, string) (net.Conn, error) { + return func(ctx context.Context, _ string) (net.Conn, error) { + return lis.DialContext(ctx) + } +} + +// dialNoClientCert connects to the test server without presenting a client cert. +// Used for the Enroll RPC (no mTLS cert yet). +func dialNoClientCert(t *testing.T, lis *bufconn.Listener) *grpc.ClientConn { + t.Helper() + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(bufDialer(lis)), + grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ + InsecureSkipVerify: true, //nolint:gosec // test-only + })), + ) + if err != nil { + t.Fatalf("dialNoClientCert: %v", err) + } + t.Cleanup(func() { conn.Close() }) + return conn +} + +// dialWithNodeCert connects to the test server presenting certPEM / keyPEM as +// the client certificate. Used after Enroll to authenticate as the enrolled node. +func dialWithNodeCert(t *testing.T, lis *bufconn.Listener, certPEM, keyPEM []byte) *grpc.ClientConn { + t.Helper() + nodeCert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + t.Fatalf("X509KeyPair: %v", err) + } + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(bufDialer(lis)), + grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ + Certificates: []tls.Certificate{nodeCert}, + InsecureSkipVerify: true, //nolint:gosec // test-only + })), + ) + if err != nil { + t.Fatalf("dialWithNodeCert: %v", err) + } + t.Cleanup(func() { conn.Close() }) + return conn +} + +// eventually polls fn until it returns true or timeout elapses. +func eventually(t *testing.T, timeout time.Duration, fn func() bool, msg string) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if fn() { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("timeout waiting for: %s", msg) +} + +// ─── Enroll tests ───────────────────────────────────────────────────────────── + +// TestEnroll_Valid exercises the happy path: issue a bootstrap token, submit a +// valid CSR, and assert the response contains the node UUID and a valid cert. +func TestEnroll_Valid(t *testing.T) { + const testNodeUUID = "test-node-enroll-valid" + b := newTestServer(t, 1, testNodeUUID) + ctx := context.Background() + + token, err := b.tokens.IssueToken(ctx, testNodeUUID) + if err != nil { + t.Fatalf("IssueToken: %v", err) + } + + keyPEM, csrPEM := genKeyAndCSR(t) + _ = keyPEM + + conn := dialNoClientCert(t, b.lis) + client := agentv1.NewAgentServiceClient(conn) + + resp, err := client.Enroll(ctx, &agentv1.EnrollRequest{ + BootstrapToken: token, + CSRPEM: csrPEM, + AgentVersion: "test-1.0", + }) + if err != nil { + t.Fatalf("Enroll: %v", err) + } + if resp.NodeUUID != testNodeUUID { + t.Errorf("NodeUUID=%q, want %q", resp.NodeUUID, testNodeUUID) + } + if len(resp.CertPEM) == 0 { + t.Error("CertPEM is empty") + } + if len(resp.CAPEM) == 0 { + t.Error("CAPEM is empty") + } + if resp.NotAfterUnix == 0 { + t.Error("NotAfterUnix is 0") + } + + // Verify the returned cert has the correct CN. + block, _ := pem.Decode(resp.CertPEM) + if block == nil { + t.Fatal("CertPEM has no PEM block") + } + cert, err := x509.ParseCertificate(block.Bytes) + if err != nil { + t.Fatalf("parse cert: %v", err) + } + if cert.Subject.CommonName != testNodeUUID { + t.Errorf("cert CN=%q, want %q", cert.Subject.CommonName, testNodeUUID) + } +} + +// TestEnroll_BadToken verifies that an invalid bootstrap token returns +// Unauthenticated and that a second call with the same token also fails. +func TestEnroll_BadToken(t *testing.T) { + const testNodeUUID = "test-node-enroll-badtoken" + b := newTestServer(t, 1, testNodeUUID) + ctx := context.Background() + + _, csrPEM := genKeyAndCSR(t) + conn := dialNoClientCert(t, b.lis) + client := agentv1.NewAgentServiceClient(conn) + + _, err := client.Enroll(ctx, &agentv1.EnrollRequest{ + BootstrapToken: "wrong-token-xxxxxxxx", + CSRPEM: csrPEM, + }) + if err == nil { + t.Fatal("expected Unauthenticated error, got nil") + } + if st, _ := status.FromError(err); st.Code() != codes.Unauthenticated { + t.Errorf("error code = %v, want Unauthenticated", st.Code()) + } +} + +// TestEnroll_TokenOneShot verifies that a bootstrap token can only be used once. +func TestEnroll_TokenOneShot(t *testing.T) { + const testNodeUUID = "test-node-oneshot" + b := newTestServer(t, 1, testNodeUUID) + ctx := context.Background() + + token, _ := b.tokens.IssueToken(ctx, testNodeUUID) + keyPEM, csrPEM := genKeyAndCSR(t) + _ = keyPEM + + conn := dialNoClientCert(t, b.lis) + client := agentv1.NewAgentServiceClient(conn) + + // First call succeeds. + if _, err := client.Enroll(ctx, &agentv1.EnrollRequest{ + BootstrapToken: token, CSRPEM: csrPEM, + }); err != nil { + t.Fatalf("first Enroll: %v", err) + } + + // Second call must fail (token consumed). + _, err := client.Enroll(ctx, &agentv1.EnrollRequest{ + BootstrapToken: token, CSRPEM: csrPEM, + }) + if err == nil { + t.Fatal("second Enroll: expected error, got nil") + } + if st, _ := status.FromError(err); st.Code() != codes.Unauthenticated { + t.Errorf("second Enroll code = %v, want Unauthenticated", st.Code()) + } +} + +// ─── helpers: full flow setup ───────────────────────────────────────────────── + +// enrollNode runs the full Enroll flow and returns the gRPC connection that +// presents the issued node cert (ready for Register/Heartbeat/Subscribe/Ack). +func enrollNode(t *testing.T, b *testServerBundle, nodeUUID string) (certPEM, keyPEM []byte, authedConn *grpc.ClientConn) { + t.Helper() + ctx := context.Background() + + token, err := b.tokens.IssueToken(ctx, nodeUUID) + if err != nil { + t.Fatalf("IssueToken: %v", err) + } + + keyPEM, csrPEM := genKeyAndCSR(t) + conn := dialNoClientCert(t, b.lis) + client := agentv1.NewAgentServiceClient(conn) + + resp, err := client.Enroll(ctx, &agentv1.EnrollRequest{ + BootstrapToken: token, + CSRPEM: csrPEM, + }) + if err != nil { + t.Fatalf("Enroll: %v", err) + } + conn.Close() + + authedConn = dialWithNodeCert(t, b.lis, resp.CertPEM, keyPEM) + return resp.CertPEM, keyPEM, authedConn +} + +// ─── Register tests ─────────────────────────────────────────────────────────── + +// TestRegister_OK verifies Register returns a valid ConfigSnapshot for a known node. +func TestRegister_OK(t *testing.T) { + const nodeUUID = "test-node-register" + b := newTestServer(t, 7, nodeUUID) + ctx := context.Background() + + _, _, conn := enrollNode(t, b, nodeUUID) + client := agentv1.NewAgentServiceClient(conn) + + snap, err := client.Register(ctx, &agentv1.RegisterRequest{ + NodeUUID: nodeUUID, + AgentVersion: "test-1.0", + }) + if err != nil { + t.Fatalf("Register: %v", err) + } + if snap.ConfigVersion != 7 { + t.Errorf("ConfigVersion=%d, want 7", snap.ConfigVersion) + } + // The mock store returns a node with RealityPBK set. + if snap.Reality == nil { + t.Error("Reality is nil, want non-nil") + } +} + +// ─── Heartbeat tests (includes Redis node:load assertion) ──────────────────── + +// TestHeartbeat_LoadWritten verifies that a Heartbeat call writes to the Redis +// node:load:{nodeUUID} hash with the reported metrics. +func TestHeartbeat_LoadWritten(t *testing.T) { + const nodeUUID = "test-node-heartbeat" + b := newTestServer(t, 5, nodeUUID) + ctx := context.Background() + + _, _, conn := enrollNode(t, b, nodeUUID) + client := agentv1.NewAgentServiceClient(conn) + + // Register first (as the agent would do in practice). + if _, err := client.Register(ctx, &agentv1.RegisterRequest{ + NodeUUID: nodeUUID, AgentVersion: "1.0", + }); err != nil { + t.Fatalf("Register: %v", err) + } + + // Send a heartbeat with known metrics. + resp, err := client.Heartbeat(ctx, &agentv1.HeartbeatRequest{ + NodeUUID: nodeUUID, + ConfigVersion: 5, + OnlinePeers: 42, + BandwidthUpBps: 1_000_000, + BandwidthDownBps: 5_000_000, + CPUPercent: 12.5, + TimestampUnix: time.Now().Unix(), + }) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if resp.ServerTimeUnix == 0 { + t.Error("ServerTimeUnix is 0") + } + if resp.NeedFullResync { + t.Error("NeedFullResync=true for matching config version, want false") + } + + // Assert Redis hash was written. + eventually(t, time.Second, func() bool { + l, ok, _ := b.load.Get(ctx, nodeUUID) + return ok && l != nil && l.OnlineCount == 42 + }, "node:load hash to contain online_count=42") + + l, _, _ := b.load.Get(ctx, nodeUUID) + if l.BandwidthUpBps != 1_000_000 { + t.Errorf("bw_up=%d, want 1000000", l.BandwidthUpBps) + } + if l.BandwidthDownBps != 5_000_000 { + t.Errorf("bw_down=%d, want 5000000", l.BandwidthDownBps) + } +} + +// TestHeartbeat_ResyncFlag verifies that a Heartbeat with a stale config_version +// triggers NeedFullResync = true. +func TestHeartbeat_ResyncFlag(t *testing.T) { + const nodeUUID = "test-node-resync" + b := newTestServer(t, 10, nodeUUID) // server is at config version 10 + ctx := context.Background() + + _, _, conn := enrollNode(t, b, nodeUUID) + client := agentv1.NewAgentServiceClient(conn) + + if _, err := client.Register(ctx, &agentv1.RegisterRequest{NodeUUID: nodeUUID}); err != nil { + t.Fatalf("Register: %v", err) + } + + resp, err := client.Heartbeat(ctx, &agentv1.HeartbeatRequest{ + NodeUUID: nodeUUID, + ConfigVersion: 5, // stale: server is at 10 + TimestampUnix: time.Now().Unix(), + }) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if !resp.NeedFullResync { + t.Error("NeedFullResync=false for stale config_version, want true") + } +} + +// ─── Subscribe / Ack tests ──────────────────────────────────────────────────── + +// TestSubscribe_PushAck exercises the full Subscribe→Push→receive→Ack cycle and +// verifies the command is removed from the Redis queue after Ack. +func TestSubscribe_PushAck(t *testing.T) { + const nodeUUID = "test-node-subscribe" + b := newTestServer(t, 1, nodeUUID) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, _, conn := enrollNode(t, b, nodeUUID) + client := agentv1.NewAgentServiceClient(conn) + + if _, err := client.Register(ctx, &agentv1.RegisterRequest{NodeUUID: nodeUUID}); err != nil { + t.Fatalf("Register: %v", err) + } + + // Open Subscribe stream. + stream, err := client.Subscribe(ctx, &agentv1.SubscribeRequest{ + NodeUUID: nodeUUID, + LastCommandID: 0, + }) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + + // Push a command via the hub. + pushed := &agentv1.Command{ + Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-sub"}}, + } + if err := b.hub.Push(ctx, nodeUUID, pushed); err != nil { + t.Fatalf("hub.Push: %v", err) + } + if pushed.CommandID == 0 { + t.Fatal("pushed.CommandID == 0: hub did not assign an ID") + } + + // Receive the command. + received, err := stream.Recv() + if err != nil { + t.Fatalf("stream.Recv: %v", err) + } + if received.CommandID != pushed.CommandID { + t.Errorf("received command_id=%d, want %d", received.CommandID, pushed.CommandID) + } + if received.Upsert.Credential.DpUUID != "dp-sub" { + t.Errorf("dp_uuid=%q, want dp-sub", received.Upsert.Credential.DpUUID) + } + + // Ack the command. + if _, err := client.Ack(ctx, &agentv1.AckRequest{ + NodeUUID: nodeUUID, + CommandID: received.CommandID, + }); err != nil { + t.Fatalf("Ack: %v", err) + } + + // Verify the command was removed from the Redis queue. + eventually(t, time.Second, func() bool { + cmds, _ := b.hub.Replay(context.Background(), nodeUUID, 0) + return len(cmds) == 0 + }, "command queue to be empty after Ack") +} + +// TestSubscribe_ReconnectResume verifies the at-least-once resume invariant: +// 1. Push commands A and B; agent receives and acks A. +// 2. Drop the stream before ack of B. +// 3. Push command C while offline. +// 4. Reconnect with last_command_id = A.CommandID. +// 5. Assert: B and C are replayed; A is NOT replayed. +func TestSubscribe_ReconnectResume(t *testing.T) { + const nodeUUID = "test-node-reconnect" + b := newTestServer(t, 1, nodeUUID) + bgCtx := context.Background() + + // Helper: enroll and return a connection with the node cert. + enrollRaw := func() (certPEM, keyPEM []byte) { + tok, err := b.tokens.IssueToken(bgCtx, nodeUUID) + if err != nil { + t.Fatalf("IssueToken: %v", err) + } + kp, csr := genKeyAndCSR(t) + c0 := dialNoClientCert(t, b.lis) + resp, err := agentv1.NewAgentServiceClient(c0).Enroll(bgCtx, + &agentv1.EnrollRequest{BootstrapToken: tok, CSRPEM: csr}) + if err != nil { + t.Fatalf("Enroll: %v", err) + } + c0.Close() + return resp.CertPEM, kp + } + + cert1, key1 := enrollRaw() + + // ─── session 1 ───────────────────────────────────────────────────────── + ctx1, cancel1 := context.WithCancel(bgCtx) + conn1 := dialWithNodeCert(t, b.lis, cert1, key1) + defer conn1.Close() + client1 := agentv1.NewAgentServiceClient(conn1) + + if _, err := client1.Register(ctx1, &agentv1.RegisterRequest{NodeUUID: nodeUUID}); err != nil { + t.Fatalf("Register session1: %v", err) + } + + stream1, err := client1.Subscribe(ctx1, &agentv1.SubscribeRequest{ + NodeUUID: nodeUUID, + LastCommandID: 0, + }) + if err != nil { + t.Fatalf("Subscribe session1: %v", err) + } + + // Push command A. + cmdA := &agentv1.Command{ + Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-A"}}, + } + if err := b.hub.Push(bgCtx, nodeUUID, cmdA); err != nil { + t.Fatalf("push A: %v", err) + } + + // Receive A. + recvA, err := stream1.Recv() + if err != nil { + t.Fatalf("recv A: %v", err) + } + if recvA.Upsert.Credential.DpUUID != "dp-A" { + t.Errorf("recv A dp_uuid=%q, want dp-A", recvA.Upsert.Credential.DpUUID) + } + + // Ack A. + if _, err := client1.Ack(bgCtx, &agentv1.AckRequest{ + NodeUUID: nodeUUID, CommandID: recvA.CommandID, + }); err != nil { + t.Fatalf("ack A: %v", err) + } + + // Push command B. + cmdB := &agentv1.Command{ + Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-B"}}, + } + if err := b.hub.Push(bgCtx, nodeUUID, cmdB); err != nil { + t.Fatalf("push B: %v", err) + } + + // Receive B (but don't ack — simulate drop before ack). + recvB, err := stream1.Recv() + if err != nil { + t.Fatalf("recv B: %v", err) + } + if recvB.Upsert.Credential.DpUUID != "dp-B" { + t.Errorf("recv B dp_uuid=%q, want dp-B", recvB.Upsert.Credential.DpUUID) + } + + // Drop stream by cancelling session 1's context. + cancel1() + time.Sleep(50 * time.Millisecond) // let server observe the cancellation + + // Push command C while node is "offline" (no active stream). + cmdC := &agentv1.Command{ + Type: agentv1.CommandTypeRevoke, + Revoke: &agentv1.RevokePayload{DpUUID: "dp-A"}, + } + if err := b.hub.Push(bgCtx, nodeUUID, cmdC); err != nil { + t.Fatalf("push C: %v", err) + } + + // ─── session 2: reconnect ────────────────────────────────────────────── + // Enroll a second time with a fresh cert (each cert is one-time use for enroll). + cert2, key2 := enrollRaw() + + ctx2, cancel2 := context.WithCancel(bgCtx) + defer cancel2() + conn2 := dialWithNodeCert(t, b.lis, cert2, key2) + defer conn2.Close() + client2 := agentv1.NewAgentServiceClient(conn2) + + if _, err := client2.Register(ctx2, &agentv1.RegisterRequest{NodeUUID: nodeUUID}); err != nil { + t.Fatalf("Register session2: %v", err) + } + + // Reconnect with last_command_id = A.CommandID (A was acked, B/C were not). + stream2, err := client2.Subscribe(ctx2, &agentv1.SubscribeRequest{ + NodeUUID: nodeUUID, + LastCommandID: recvA.CommandID, + }) + if err != nil { + t.Fatalf("Subscribe session2: %v", err) + } + + // Collect resumed commands (expect exactly B and C). + type recvResult struct { + cmd *agentv1.Command + err error + } + recvCh := make(chan recvResult, 3) + go func() { + for { + cmd, err := stream2.Recv() + recvCh <- recvResult{cmd, err} + if err != nil { + return + } + } + }() + + var resumed []*agentv1.Command + timeout := time.After(3 * time.Second) + for len(resumed) < 2 { + select { + case r := <-recvCh: + if r.err != nil { + t.Fatalf("stream2.Recv: %v", r.err) + } + resumed = append(resumed, r.cmd) + case <-timeout: + t.Fatalf("timeout: got %d resumed commands, want 2", len(resumed)) + } + } + + // Resumed commands must be B and C in order. + if resumed[0].Upsert.Credential.DpUUID != "dp-B" { + t.Errorf("resumed[0] dp_uuid=%q, want dp-B", resumed[0].Upsert.Credential.DpUUID) + } + if resumed[1].Revoke.DpUUID != "dp-A" { + t.Errorf("resumed[1] revoke dp_uuid=%q, want dp-A", resumed[1].Revoke.DpUUID) + } + + // Command A must NOT be in the resumed batch (it was acked). + for _, c := range resumed { + if c.CommandID == recvA.CommandID { + t.Errorf("command A (id=%d) was re-delivered, but it was acked", recvA.CommandID) + } + } +} + +// ─── ReportUsage test ───────────────────────────────────────────────────────── + +// TestReportUsage_Accumulates verifies that ReportUsage resolves dp_uuid → user_id +// and records usage in the mock store. +func TestReportUsage_Accumulates(t *testing.T) { + const nodeUUID = "test-node-usage" + b := newTestServer(t, 1, nodeUUID) + ctx := context.Background() + + _, _, conn := enrollNode(t, b, nodeUUID) + client := agentv1.NewAgentServiceClient(conn) + + if _, err := client.Register(ctx, &agentv1.RegisterRequest{NodeUUID: nodeUUID}); err != nil { + t.Fatalf("Register: %v", err) + } + + now := time.Now() + _, err := client.ReportUsage(ctx, &agentv1.UsageReport{ + NodeUUID: nodeUUID, + WindowStartUnix: now.Add(-time.Minute).Unix(), + WindowEndUnix: now.Unix(), + Entries: []*agentv1.UsageEntry{ + {DpUUID: "dp-user1", BytesUp: 100, BytesDown: 200, SessionMinutes: 3}, + {DpUUID: "dp-unknown", BytesUp: 50, BytesDown: 50, SessionMinutes: 1}, // skipped + }, + }) + if err != nil { + t.Fatalf("ReportUsage: %v", err) + } + + log := b.store.usageLog() + if len(log) != 1 { + t.Fatalf("usage log has %d entries, want 1 (dp-unknown skipped)", len(log)) + } + if log[0].UserID != 101 { + t.Errorf("user_id=%d, want 101", log[0].UserID) + } + if log[0].BytesUp != 100 || log[0].BytesDown != 200 { + t.Errorf("bytes_up=%d bytes_down=%d, want 100/200", + log[0].BytesUp, log[0].BytesDown) + } +} diff --git a/server/internal/nodes/handler_grpc.go b/server/internal/nodes/handler_grpc.go new file mode 100644 index 0000000..fd203d9 --- /dev/null +++ b/server/internal/nodes/handler_grpc.go @@ -0,0 +1,265 @@ +package nodes + +import ( + "context" + "log/slog" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/wangjia/pangolin/server/internal/mtls" + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +// Handler implements agentv1.AgentServiceServer. +// Construct via NewHandler after wiring the individual components. +type Handler struct { + ca *mtls.CA + tokens *mtls.BootstrapTokenManager + hub *Hub + store NodeStore + load *LoadCache +} + +// NewHandler creates a Handler with all dependencies injected. +func NewHandler( + ca *mtls.CA, + tokens *mtls.BootstrapTokenManager, + hub *Hub, + store NodeStore, + load *LoadCache, +) *Handler { + return &Handler{ + ca: ca, + tokens: tokens, + hub: hub, + store: store, + load: load, + } +} + +// Enroll is the one-time node enrollment RPC. +// +// Flow: consume bootstrap token (one-shot, returns Unauthenticated if invalid) → +// sign CSR with CN = nodeUUID → return cert + CA cert. +// The agent generates the key on-node so the private key never leaves it. +func (h *Handler) Enroll(ctx context.Context, req *agentv1.EnrollRequest) (*agentv1.EnrollResponse, error) { + if req.BootstrapToken == "" { + return nil, status.Error(codes.Unauthenticated, "missing bootstrap_token") + } + if len(req.CSRPEM) == 0 { + return nil, status.Error(codes.InvalidArgument, "missing csr_pem") + } + + // ConsumeToken is atomic GETDEL: second call always fails. + nodeUUID, err := h.tokens.ConsumeToken(ctx, req.BootstrapToken) + if err != nil { + return nil, status.Errorf(codes.Unauthenticated, "invalid bootstrap token: %v", err) + } + + certPEM, err := h.ca.SignCSR(req.CSRPEM, nodeUUID) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "sign CSR: %v", err) + } + + return &agentv1.EnrollResponse{ + NodeUUID: nodeUUID, + CertPEM: certPEM, + CAPEM: h.ca.CAPEM(), + NotAfterUnix: time.Now().Add(90 * 24 * time.Hour).Unix(), + }, nil +} + +// Register returns the current configuration snapshot for the caller node. +// +// The node UUID is taken from the verified mTLS certificate CN (injected by +// mtls.UnaryServerInterceptor); the req.NodeUUID field is informational only. +// Returns NotFound when no nodes table row exists for the UUID. +func (h *Handler) Register(ctx context.Context, req *agentv1.RegisterRequest) (*agentv1.ConfigSnapshot, error) { + nodeUUID, ok := mtls.NodeUUIDFromContext(ctx) + if !ok { + return nil, status.Error(codes.Unauthenticated, "missing node identity in context") + } + + node, err := h.store.NodeByUUID(ctx, nodeUUID) + if err != nil { + return nil, status.Errorf(codes.Internal, "node lookup: %v", err) + } + if node == nil { + return nil, status.Errorf(codes.NotFound, "node %q not in node catalogue", nodeUUID) + } + + configVer, err := h.store.ConfigVersion(ctx) + if err != nil { + return nil, status.Errorf(codes.Internal, "config version: %v", err) + } + + creds, err := h.store.CredentialsForNode(ctx, nodeUUID) + if err != nil { + return nil, status.Errorf(codes.Internal, "credentials: %v", err) + } + + snap := &agentv1.ConfigSnapshot{ + ConfigVersion: configVer, + Credentials: creds, + } + + // Populate inbound configs from the nodes row. + if node.RealityPBK != "" { + snap.Reality = &agentv1.RealityInbound{ + PrivateKey: node.RealityPBK, + ServerName: node.RealitySNI, + } + } + if node.Hy2Port.Valid { + snap.Hy2 = &agentv1.Hy2Inbound{ + ListenPort: node.Hy2Port.Int32, + } + } + + return snap, nil +} + +// Heartbeat records the node's current load metrics and signals whether a full +// resync is required (the node's config_version is behind the server's). +func (h *Handler) Heartbeat(ctx context.Context, req *agentv1.HeartbeatRequest) (*agentv1.HeartbeatResponse, error) { + nodeUUID, ok := mtls.NodeUUIDFromContext(ctx) + if !ok { + return nil, status.Error(codes.Unauthenticated, "missing node identity in context") + } + + // Best-effort load write; a failure here must not break the heartbeat loop. + if err := h.load.Set(ctx, nodeUUID, NodeLoad{ + OnlineCount: int64(req.OnlinePeers), + BandwidthUpBps: req.BandwidthUpBps, + BandwidthDownBps: req.BandwidthDownBps, + CPUPercent: req.CPUPercent, + }); err != nil { + slog.Warn("nodes.Handler.Heartbeat: load write failed", + "node_uuid", nodeUUID, "err", err) + } + + configVer, err := h.store.ConfigVersion(ctx) + if err != nil { + return nil, status.Errorf(codes.Internal, "config version: %v", err) + } + + // Request full resync only when the node reports a positive version that is + // strictly behind the server (req.ConfigVersion == 0 means "freshly registered"). + needResync := req.ConfigVersion > 0 && req.ConfigVersion < configVer + + return &agentv1.HeartbeatResponse{ + NeedFullResync: needResync, + ServerTimeUnix: time.Now().Unix(), + }, nil +} + +// Subscribe streams pending and future commands to the caller node. +// +// On each call the handler: +// 1. Registers with the hub (before replaying to avoid a delivery gap). +// 2. Replays all unacked commands with id > req.LastCommandID from the +// Redis queue — ensuring continuity after a reconnect. +// 3. Pumps newly arriving commands until the stream context is cancelled. +func (h *Handler) Subscribe(req *agentv1.SubscribeRequest, stream agentv1.AgentService_SubscribeServer) error { + ctx := stream.Context() + nodeUUID, ok := mtls.NodeUUIDFromContext(ctx) + if !ok { + return status.Error(codes.Unauthenticated, "missing node identity in context") + } + + // Register BEFORE replaying to avoid losing commands that arrive in between. + ch, done := h.hub.Register(nodeUUID) + defer done() + + // Replay unacked backlog (commands in queue with id > last_command_id). + backlog, err := h.hub.Replay(ctx, nodeUUID, req.LastCommandID) + if err != nil { + return status.Errorf(codes.Internal, "replay: %v", err) + } + + // Track the highest command_id sent from the backlog so the pump loop can + // skip duplicates. A command pushed between Register and Replay ends up in + // both the Redis queue (visible to Replay) and the local channel (delivered + // by Push); without this guard the client would receive it twice. + var maxReplayedID int64 + for _, cmd := range backlog { + if err := stream.Send(cmd); err != nil { + return err + } + if cmd.CommandID > maxReplayedID { + maxReplayedID = cmd.CommandID + } + } + + // Pump live commands until the stream breaks (client disconnect / ctx cancel). + for { + select { + case <-ctx.Done(): + return ctx.Err() + case cmd, ok := <-ch: + if !ok { + return nil + } + // Skip commands that were already sent from the backlog replay above. + if cmd.CommandID > 0 && cmd.CommandID <= maxReplayedID { + continue + } + if err := stream.Send(cmd); err != nil { + return err + } + } + } +} + +// Ack removes commandID from the node's pending queue. +// A repeated Ack for the same commandID is a no-op. +func (h *Handler) Ack(ctx context.Context, req *agentv1.AckRequest) (*agentv1.AckResponse, error) { + nodeUUID, ok := mtls.NodeUUIDFromContext(ctx) + if !ok { + return nil, status.Error(codes.Unauthenticated, "missing node identity in context") + } + if err := h.hub.Ack(ctx, nodeUUID, req.CommandID); err != nil { + return nil, status.Errorf(codes.Internal, "ack: %v", err) + } + return &agentv1.AckResponse{}, nil +} + +// ReportUsage accumulates per-dp_uuid bandwidth and session metrics into usage_daily. +// +// The node has no knowledge of user accounts — it only sends opaque dp_uuids and +// byte/minute counters. The control plane resolves dp_uuid → user_id internally. +// Entries for unknown dp_uuids are silently skipped (user may have been deleted). +func (h *Handler) ReportUsage(ctx context.Context, req *agentv1.UsageReport) (*agentv1.UsageAck, error) { + _, ok := mtls.NodeUUIDFromContext(ctx) + if !ok { + return nil, status.Error(codes.Unauthenticated, "missing node identity in context") + } + + // Use the window end time to determine the calendar date for usage_daily. + windowEnd := time.Unix(req.WindowEndUnix, 0).UTC() + date := windowEnd.Truncate(24 * time.Hour) + + for _, entry := range req.Entries { + if entry.DpUUID == "" { + continue + } + userID, found, err := h.store.UserIDByDpUUID(ctx, entry.DpUUID) + if err != nil { + slog.Warn("nodes.Handler.ReportUsage: dp_uuid lookup failed", + "dp_uuid", entry.DpUUID, "err", err) + continue // best-effort; don't fail the whole report + } + if !found { + continue + } + if err := h.store.AccumulateUsage(ctx, userID, date, + entry.BytesUp, entry.BytesDown, entry.SessionMinutes, + ); err != nil { + slog.Warn("nodes.Handler.ReportUsage: accumulate failed", + "user_id", userID, "err", err) + } + } + return &agentv1.UsageAck{}, nil +} diff --git a/server/internal/nodes/hub.go b/server/internal/nodes/hub.go new file mode 100644 index 0000000..8affc1c --- /dev/null +++ b/server/internal/nodes/hub.go @@ -0,0 +1,228 @@ +package nodes + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "strconv" + "sync" + + "github.com/redis/go-redis/v9" + + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +const ( + // cmdQueuePrefix is the Redis ZSET key for each node's pending command queue. + // Score = command_id (float64 cast of int64); member = JSON-encoded Command. + // Commands remain in the queue until Ack'd; reconnecting nodes replay from afterID. + cmdQueuePrefix = "node:cmdq:" + + // cmdPubChannel is the Redis Pub/Sub channel for cross-instance command delivery. + // Published whenever a command cannot be delivered locally (node connected elsewhere). + cmdPubChannel = "node:cmd" + + // cmdIDKey is the Redis counter for auto-generating monotonically increasing command IDs. + cmdIDKey = "pangolin:node:next_cmdid" + + // subChanCap is the size of each local subscriber's delivery channel buffer. + subChanCap = 256 +) + +// pubSubPayload is the message body published to cmdPubChannel. +type pubSubPayload struct { + NodeUUID string `json:"n"` + Cmd *agentv1.Command `json:"c"` +} + +// Hub routes commands to connected nodes. +// +// Persistence: every pushed command is stored in a per-node Redis ZSET +// (node:cmdq:{uuid}) so the queue survives instance restarts. +// +// Local delivery: each active Subscribe stream holds a buffered channel. +// When a node is connected locally, commands are sent directly to that channel. +// +// Cross-instance delivery: when the target node is not connected to this instance, +// the command is published on cmdPubChannel so any other instance holding the node's +// stream can pick it up and forward it. +// +// At-least-once semantics: commands are removed from the queue only on Ack. +// A reconnecting node replays from its last acknowledged command_id. +type Hub struct { + mu sync.RWMutex + subs map[string]chan *agentv1.Command + + rdb *redis.Client +} + +// NewHub creates a Hub backed by rdb. +// Call Start(ctx) to activate the cross-instance pub/sub goroutine. +func NewHub(rdb *redis.Client) *Hub { + return &Hub{ + subs: make(map[string]chan *agentv1.Command), + rdb: rdb, + } +} + +// Start begins the background pub/sub subscriber goroutine. +// It runs until ctx is cancelled. +func (h *Hub) Start(ctx context.Context) { + go h.runPubSub(ctx) +} + +func (h *Hub) runPubSub(ctx context.Context) { + sub := h.rdb.Subscribe(ctx, cmdPubChannel) + defer func() { _ = sub.Close() }() + + ch := sub.Channel() + for { + select { + case <-ctx.Done(): + return + case msg, ok := <-ch: + if !ok { + return + } + var p pubSubPayload + if err := json.Unmarshal([]byte(msg.Payload), &p); err != nil { + slog.Warn("nodes.hub: bad pub/sub message", "err", err) + continue + } + h.deliverLocal(p.NodeUUID, p.Cmd) + } + } +} + +// deliverLocal sends cmd to the local subscriber for nodeUUID (non-blocking). +// Returns true if the command was placed on the channel. +func (h *Hub) deliverLocal(nodeUUID string, cmd *agentv1.Command) bool { + h.mu.RLock() + ch, ok := h.subs[nodeUUID] + h.mu.RUnlock() + if !ok { + return false + } + select { + case ch <- cmd: + return true + default: + slog.Warn("nodes.hub: subscriber channel full, dropping command", + "node_uuid", nodeUUID, "command_id", cmd.CommandID) + return false + } +} + +// Register adds a local subscriber for nodeUUID and returns: +// - ch: the receive channel on which live commands will arrive. +// - done: cleanup function that must be called when the stream ends. +// +// Register must be called before Replay to avoid missing commands that arrive +// between the two operations. +func (h *Hub) Register(nodeUUID string) (<-chan *agentv1.Command, func()) { + ch := make(chan *agentv1.Command, subChanCap) + h.mu.Lock() + h.subs[nodeUUID] = ch + h.mu.Unlock() + return ch, func() { + h.mu.Lock() + // Guard against re-registration: only delete if this is still the active channel. + if h.subs[nodeUUID] == ch { + delete(h.subs, nodeUUID) + } + h.mu.Unlock() + } +} + +// Replay returns all unacked commands for nodeUUID with command_id > afterID, ordered ascending. +// This is called during Subscribe to re-send commands the agent may have missed. +func (h *Hub) Replay(ctx context.Context, nodeUUID string, afterID int64) ([]*agentv1.Command, error) { + key := cmdQueuePrefix + nodeUUID + min := fmt.Sprintf("(%d", afterID) // exclusive: score > afterID + strs, err := h.rdb.ZRangeByScore(ctx, key, &redis.ZRangeBy{ + Min: min, + Max: "+inf", + }).Result() + if err != nil { + return nil, fmt.Errorf("nodes.hub.Replay: %w", err) + } + cmds := make([]*agentv1.Command, 0, len(strs)) + for _, s := range strs { + var c agentv1.Command + if err := json.Unmarshal([]byte(s), &c); err != nil { + slog.Warn("nodes.hub: corrupt command in queue, skipping", + "err", err, "node_uuid", nodeUUID) + continue + } + cmds = append(cmds, &c) + } + return cmds, nil +} + +// Push enqueues cmd for nodeUUID. +// If cmd.CommandID == 0, a new monotonic ID is assigned via Redis INCR. +// The command is persisted in the node's Redis queue then delivered: +// - directly to the local subscriber channel if this instance holds the stream, +// - otherwise published on cmdPubChannel for cross-instance delivery. +func (h *Hub) Push(ctx context.Context, nodeUUID string, cmd *agentv1.Command) error { + // Assign ID if not set. + if cmd.CommandID == 0 { + id, err := h.rdb.Incr(ctx, cmdIDKey).Result() + if err != nil { + return fmt.Errorf("nodes.hub.Push: generate id: %w", err) + } + cmd.CommandID = id + } + + // Persist to Redis ZSET (durable, survives instance restart). + data, err := json.Marshal(cmd) + if err != nil { + return fmt.Errorf("nodes.hub.Push: marshal: %w", err) + } + if err := h.rdb.ZAdd(ctx, cmdQueuePrefix+nodeUUID, redis.Z{ + Score: float64(cmd.CommandID), + Member: string(data), + }).Err(); err != nil { + return fmt.Errorf("nodes.hub.Push: zadd: %w", err) + } + + // Try local delivery first (fastest path, no pub/sub latency). + if h.deliverLocal(nodeUUID, cmd) { + return nil + } + + // Publish for cross-instance delivery. If the node is not connected anywhere, + // this is a no-op and the node will pick up the command on next reconnect via Replay. + payload, err := json.Marshal(pubSubPayload{NodeUUID: nodeUUID, Cmd: cmd}) + if err != nil { + return fmt.Errorf("nodes.hub.Push: marshal pubsub: %w", err) + } + if err := h.rdb.Publish(ctx, cmdPubChannel, payload).Err(); err != nil { + // Non-fatal: command is already durable. + slog.Warn("nodes.hub: publish failed (command still queued)", + "node_uuid", nodeUUID, "command_id", cmd.CommandID, "err", err) + } + return nil +} + +// Broadcast pushes cmd to each nodeUUID with a freshly assigned command_id per node. +// A Broadcast cmd should be passed with CommandID = 0. +func (h *Hub) Broadcast(ctx context.Context, nodeUUIDs []string, cmd *agentv1.Command) error { + for _, uuid := range nodeUUIDs { + clone := *cmd + clone.CommandID = 0 // each push gets an independent ID + if err := h.Push(ctx, uuid, &clone); err != nil { + slog.Warn("nodes.hub.Broadcast: push failed", + "node_uuid", uuid, "err", err) + } + } + return nil +} + +// Ack removes commandID from the pending queue for nodeUUID. +// A second Ack for the same commandID is a no-op (idempotent). +func (h *Hub) Ack(ctx context.Context, nodeUUID string, commandID int64) error { + score := strconv.FormatInt(commandID, 10) + return h.rdb.ZRemRangeByScore(ctx, cmdQueuePrefix+nodeUUID, score, score).Err() +} diff --git a/server/internal/nodes/hub_test.go b/server/internal/nodes/hub_test.go new file mode 100644 index 0000000..6ddc3e1 --- /dev/null +++ b/server/internal/nodes/hub_test.go @@ -0,0 +1,386 @@ +package nodes_test + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + + "github.com/wangjia/pangolin/server/internal/nodes" + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +// newTestHub creates a Hub backed by a fresh miniredis instance. +func newTestHub(t *testing.T) (*nodes.Hub, *redis.Client, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { rdb.Close() }) + return nodes.NewHub(rdb), rdb, mr +} + +// ─── Push / Replay / Ack ───────────────────────────────────────────────────── + +// TestHub_PushReplay verifies that pushed commands are retrievable via Replay, +// filtered by afterID (exclusive), and removed from the queue on Ack. +func TestHub_PushReplay(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + const nodeUUID = "node-replay-test" + + // Push three commands with explicit IDs so the test is deterministic. + cmds := []*agentv1.Command{ + {CommandID: 10, Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-a"}}}, + {CommandID: 20, Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-b"}}}, + {CommandID: 30, Type: agentv1.CommandTypeRevoke, + Revoke: &agentv1.RevokePayload{DpUUID: "dp-a"}}, + } + for _, c := range cmds { + if err := hub.Push(ctx, nodeUUID, c); err != nil { + t.Fatalf("Push id=%d: %v", c.CommandID, err) + } + } + + // Replay from after 10 → must return 20 and 30 only. + got, err := hub.Replay(ctx, nodeUUID, 10) + if err != nil { + t.Fatalf("Replay: %v", err) + } + if len(got) != 2 { + t.Fatalf("Replay returned %d commands, want 2", len(got)) + } + if got[0].CommandID != 20 || got[1].CommandID != 30 { + t.Errorf("Replay ids = [%d %d], want [20 30]", + got[0].CommandID, got[1].CommandID) + } + + // Ack command 20. + if err := hub.Ack(ctx, nodeUUID, 20); err != nil { + t.Fatalf("Ack 20: %v", err) + } + + // Replay from 0: 10 and 30 remain; 20 was removed. + got, err = hub.Replay(ctx, nodeUUID, 0) + if err != nil { + t.Fatalf("Replay after ack: %v", err) + } + if len(got) != 2 { + t.Fatalf("after Ack 20: Replay returned %d commands, want 2", len(got)) + } + ids := []int64{got[0].CommandID, got[1].CommandID} + if ids[0] != 10 || ids[1] != 30 { + t.Errorf("after Ack 20: ids = %v, want [10 30]", ids) + } +} + +// TestHub_AutoID verifies that Hub assigns monotonically increasing IDs when +// cmd.CommandID == 0. +func TestHub_AutoID(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + const nodeUUID = "node-auto-id" + + c1 := &agentv1.Command{Type: agentv1.CommandTypeLifecycle, + Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain}} + c2 := &agentv1.Command{Type: agentv1.CommandTypeLifecycle, + Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionResume}} + + if err := hub.Push(ctx, nodeUUID, c1); err != nil { + t.Fatalf("Push c1: %v", err) + } + if err := hub.Push(ctx, nodeUUID, c2); err != nil { + t.Fatalf("Push c2: %v", err) + } + + if c1.CommandID == 0 { + t.Error("c1.CommandID still 0 after Push") + } + if c2.CommandID == 0 { + t.Error("c2.CommandID still 0 after Push") + } + if c1.CommandID >= c2.CommandID { + t.Errorf("c1.CommandID=%d >= c2.CommandID=%d; want strictly increasing", + c1.CommandID, c2.CommandID) + } +} + +// TestHub_AckIdempotent verifies that a second Ack for the same command is a no-op. +func TestHub_AckIdempotent(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + + cmd := &agentv1.Command{CommandID: 99, Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "x"}}} + if err := hub.Push(ctx, "node-ack-idem", cmd); err != nil { + t.Fatalf("Push: %v", err) + } + + // First Ack: removes the command. + if err := hub.Ack(ctx, "node-ack-idem", 99); err != nil { + t.Fatalf("Ack 1: %v", err) + } + // Second Ack: must not return an error (ZREMRANGEBYSCORE with no match = 0 removed). + if err := hub.Ack(ctx, "node-ack-idem", 99); err != nil { + t.Errorf("Ack 2 (idempotent): unexpected error: %v", err) + } + + // Queue should be empty. + got, _ := hub.Replay(ctx, "node-ack-idem", 0) + if len(got) != 0 { + t.Errorf("after double Ack, queue has %d commands, want 0", len(got)) + } +} + +// ─── Local delivery ─────────────────────────────────────────────────────────── + +// TestHub_LocalDelivery verifies that a Push reaches a registered subscriber's +// channel when the node is connected to the same hub instance. +func TestHub_LocalDelivery(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + const nodeUUID = "node-local-delivery" + + ch, done := hub.Register(nodeUUID) + defer done() + + cmd := &agentv1.Command{ + Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-local"}}, + } + if err := hub.Push(ctx, nodeUUID, cmd); err != nil { + t.Fatalf("Push: %v", err) + } + + select { + case received := <-ch: + if received.Upsert.Credential.DpUUID != "dp-local" { + t.Errorf("got dp_uuid=%q, want dp-local", received.Upsert.Credential.DpUUID) + } + case <-time.After(time.Second): + t.Fatal("timeout: command not delivered to local subscriber") + } +} + +// TestHub_RegisterUnregister verifies that after done() is called, the node no +// longer receives commands on the old channel. +func TestHub_RegisterUnregister(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + const nodeUUID = "node-unreg" + + ch, done := hub.Register(nodeUUID) + done() // immediately unregister + + cmd := &agentv1.Command{ + CommandID: 42, + Type: agentv1.CommandTypeLifecycle, + Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain}, + } + if err := hub.Push(ctx, nodeUUID, cmd); err != nil { + t.Fatalf("Push: %v", err) + } + + select { + case m := <-ch: + // The channel might receive the command if it was buffered before unregister. + // That is acceptable (it was still in the buffer at send time). + // What must NOT happen is a panic or deadlock. + _ = m + case <-time.After(100 * time.Millisecond): + // Expected: no delivery since we unregistered. + } +} + +// ─── Replay after reconnect ─────────────────────────────────────────────────── + +// TestHub_ReconnectResume verifies the "only resume unacked commands" invariant: +// - Push commands 1, 2, 3 +// - Ack 1 and 2 +// - Reconnect with last_command_id = 3 (agent already has 3) +// - Push command 4 +// - Replay from 3 → only command 4 delivered +func TestHub_ReconnectResume(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + const nodeUUID = "node-reconnect" + + push := func(id int64) { + c := &agentv1.Command{ + CommandID: id, + Type: agentv1.CommandTypeLifecycle, + Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain}, + } + if err := hub.Push(ctx, nodeUUID, c); err != nil { + t.Fatalf("Push %d: %v", id, err) + } + } + + push(1) + push(2) + push(3) + + // Ack 1 and 2 (simulating a previous session that processed them). + for _, id := range []int64{1, 2} { + if err := hub.Ack(ctx, nodeUUID, id); err != nil { + t.Fatalf("Ack %d: %v", id, err) + } + } + // Command 3 was delivered but not acked (stream dropped before ack). + + // Push command 4 (arrived while node was offline). + push(4) + + // Reconnect: agent sends last_command_id = 3 (highest it processed). + // Should receive command 3 (in queue, not acked) and command 4. + got, err := hub.Replay(ctx, nodeUUID, 3) + if err != nil { + t.Fatalf("Replay: %v", err) + } + if len(got) != 1 { + t.Fatalf("Replay from 3: got %d commands, want 1 (only cmd 4)", len(got)) + } + if got[0].CommandID != 4 { + t.Errorf("Replay from 3: got command_id=%d, want 4", got[0].CommandID) + } + + // Reconnect with last_command_id = 2 (simulating agent that lost cmd 3): + // should receive cmd 3 and 4. + got, err = hub.Replay(ctx, nodeUUID, 2) + if err != nil { + t.Fatalf("Replay: %v", err) + } + if len(got) != 2 { + t.Fatalf("Replay from 2: got %d commands, want 2 (cmds 3, 4)", len(got)) + } + if got[0].CommandID != 3 || got[1].CommandID != 4 { + t.Errorf("Replay from 2: ids = [%d, %d], want [3, 4]", + got[0].CommandID, got[1].CommandID) + } +} + +// ─── Cross-instance pub/sub ─────────────────────────────────────────────────── + +// TestHub_CrossInstance starts two Hub instances connected to the same Redis, +// subscribes a node on hub1, pushes a command via hub2, and verifies delivery. +// This is the key validation that cross-instance command routing works. +func TestHub_CrossInstance(t *testing.T) { + mr := miniredis.RunT(t) + + newHub := func() *nodes.Hub { + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { rdb.Close() }) + return nodes.NewHub(rdb) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + hub1 := newHub() + hub2 := newHub() + hub1.Start(ctx) + hub2.Start(ctx) + + const nodeUUID = "node-cross-instance" + ch, done := hub1.Register(nodeUUID) + defer done() + + // Give the pub/sub goroutines a moment to subscribe to the Redis channel. + time.Sleep(50 * time.Millisecond) + + cmd := &agentv1.Command{ + Type: agentv1.CommandTypeUpsert, + Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-cross"}}, + } + if err := hub2.Push(ctx, nodeUUID, cmd); err != nil { + t.Fatalf("hub2.Push: %v", err) + } + + select { + case received := <-ch: + if received.Upsert.Credential.DpUUID != "dp-cross" { + t.Errorf("got dp_uuid=%q, want dp-cross", received.Upsert.Credential.DpUUID) + } + t.Logf("cross-instance delivery: command_id=%d", received.CommandID) + case <-time.After(3 * time.Second): + t.Fatal("timeout: cross-instance command not delivered") + } +} + +// TestHub_CrossInstance_QueuePersisted verifies that a cross-instance push +// persists the command in Redis so Replay works even without a live subscriber. +func TestHub_CrossInstance_QueuePersisted(t *testing.T) { + mr := miniredis.RunT(t) + newHub := func() *nodes.Hub { + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { rdb.Close() }) + return nodes.NewHub(rdb) + } + + ctx := context.Background() + hub1 := newHub() // not started (no pubsub), no local subscriber + hub2 := newHub() + + cmd := &agentv1.Command{ + Type: agentv1.CommandTypeRevoke, + Revoke: &agentv1.RevokePayload{DpUUID: "dp-persist"}, + } + if err := hub2.Push(ctx, "offline-node", cmd); err != nil { + t.Fatalf("hub2.Push: %v", err) + } + if cmd.CommandID == 0 { + t.Fatal("command_id was not assigned") + } + + // hub1 replays from the persistent queue (simulating reconnect). + got, err := hub1.Replay(ctx, "offline-node", 0) + if err != nil { + t.Fatalf("hub1.Replay: %v", err) + } + if len(got) != 1 { + t.Fatalf("got %d commands, want 1", len(got)) + } + if got[0].Revoke.DpUUID != "dp-persist" { + t.Errorf("got dp_uuid=%q, want dp-persist", got[0].Revoke.DpUUID) + } +} + +// TestHub_Broadcast delivers a command to multiple nodes. +func TestHub_Broadcast(t *testing.T) { + hub, _, _ := newTestHub(t) + ctx := context.Background() + + const n = 3 + nodes2 := make([]string, n) + channels := make([]<-chan *agentv1.Command, n) + dones := make([]func(), n) + + for i := 0; i < n; i++ { + uuid := "node-bcast-" + string(rune('A'+i)) + nodes2[i] = uuid + channels[i], dones[i] = hub.Register(uuid) + defer dones[i]() + } + + cmd := &agentv1.Command{ + Type: agentv1.CommandTypeLifecycle, + Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain}, + } + if err := hub.Broadcast(ctx, nodes2, cmd); err != nil { + t.Fatalf("Broadcast: %v", err) + } + + for i, ch := range channels { + select { + case received := <-ch: + if received.Type != agentv1.CommandTypeLifecycle { + t.Errorf("node %d: wrong type %v", i, received.Type) + } + case <-time.After(time.Second): + t.Errorf("node %d: timeout waiting for broadcast", i) + } + } +} diff --git a/server/internal/nodes/load.go b/server/internal/nodes/load.go new file mode 100644 index 0000000..c2af1c7 --- /dev/null +++ b/server/internal/nodes/load.go @@ -0,0 +1,83 @@ +package nodes + +import ( + "context" + "fmt" + "time" + + "github.com/redis/go-redis/v9" +) + +const ( + // nodeLoadKeyPrefix is the Redis HASH key prefix for per-node runtime load metrics. + // Full key: node:load:{node_uuid} (doc/03 §4) + nodeLoadKeyPrefix = "node:load:" + + // nodeLoadTTL is the TTL on each node:load hash. A node that stops sending + // heartbeats has its entry automatically evicted after 90 seconds. + nodeLoadTTL = 90 * time.Second +) + +// NodeLoad holds the runtime metrics reported in each heartbeat. +type NodeLoad struct { + OnlineCount int64 + BandwidthUpBps int64 + BandwidthDownBps int64 + CPUPercent float64 +} + +// LoadCache reads and writes node:load Redis hashes. +// It accepts redis.Cmdable so it can be used with a *redis.Client or pipeline. +type LoadCache struct { + rdb redis.Cmdable +} + +// NewLoadCache creates a LoadCache backed by rdb. +func NewLoadCache(rdb redis.Cmdable) *LoadCache { + return &LoadCache{rdb: rdb} +} + +// Set writes nodeUUID's current load metrics to Redis and refreshes the TTL. +// The hash fields match the names in doc/03 §4: online_count, bw_up, bw_down, cpu. +func (c *LoadCache) Set(ctx context.Context, nodeUUID string, load NodeLoad) error { + key := nodeLoadKeyPrefix + nodeUUID + pipe := c.rdb.Pipeline() + pipe.HSet(ctx, key, + "online_count", load.OnlineCount, + "bw_up", load.BandwidthUpBps, + "bw_down", load.BandwidthDownBps, + "cpu", load.CPUPercent, + ) + pipe.Expire(ctx, key, nodeLoadTTL) + if _, err := pipe.Exec(ctx); err != nil { + return fmt.Errorf("nodes.LoadCache.Set: %w", err) + } + return nil +} + +// Get reads nodeUUID's last-known load. Returns (nil, false, nil) if the key has +// expired or was never set. Fields missing from the hash default to zero. +func (c *LoadCache) Get(ctx context.Context, nodeUUID string) (*NodeLoad, bool, error) { + key := nodeLoadKeyPrefix + nodeUUID + vals, err := c.rdb.HGetAll(ctx, key).Result() + if err != nil { + return nil, false, fmt.Errorf("nodes.LoadCache.Get: %w", err) + } + if len(vals) == 0 { + return nil, false, nil + } + var l NodeLoad + if v, ok := vals["online_count"]; ok { + _, _ = fmt.Sscan(v, &l.OnlineCount) + } + if v, ok := vals["bw_up"]; ok { + _, _ = fmt.Sscan(v, &l.BandwidthUpBps) + } + if v, ok := vals["bw_down"]; ok { + _, _ = fmt.Sscan(v, &l.BandwidthDownBps) + } + if v, ok := vals["cpu"]; ok { + _, _ = fmt.Sscan(v, &l.CPUPercent) + } + return &l, true, nil +} diff --git a/server/internal/nodes/service.go b/server/internal/nodes/service.go new file mode 100644 index 0000000..4da557f --- /dev/null +++ b/server/internal/nodes/service.go @@ -0,0 +1,60 @@ +package nodes + +import ( + "github.com/redis/go-redis/v9" + + "github.com/wangjia/pangolin/server/internal/mtls" + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +// Service bundles all nodes-domain components and exposes assembly helpers. +// Construct via NewService after initialising each dependency independently. +type Service struct { + handler *Handler + hub *Hub + store NodeStore + load *LoadCache +} + +// NewService wires up the full nodes service from its dependencies. +// +// - ca / tokens / crl: from the mtls package (task 5b) +// - rdb: Redis client (for hub + load cache) +// - store: NodeStore implementation (SQLNodeStore in production; mock in tests) +func NewService( + ca *mtls.CA, + tokens *mtls.BootstrapTokenManager, + rdb *redis.Client, + store NodeStore, +) *Service { + hub := NewHub(rdb) + load := NewLoadCache(rdb) + handler := NewHandler(ca, tokens, hub, store, load) + return &Service{ + handler: handler, + hub: hub, + store: store, + load: load, + } +} + +// Handler returns the AgentServiceServer implementation for gRPC registration. +func (s *Service) Handler() agentv1.AgentServiceServer { + return s.handler +} + +// Hub exposes the command routing hub for callers (e.g. task 5d/5e) that need to +// Push or Broadcast commands. +func (s *Service) Hub() *Hub { + return s.hub +} + +// Store exposes the NodeStore for callers that need direct DB access. +func (s *Service) Store() NodeStore { + return s.store +} + +// Load exposes the LoadCache for callers that display per-node load metrics. +func (s *Service) Load() *LoadCache { + return s.load +} diff --git a/server/internal/nodes/store.go b/server/internal/nodes/store.go new file mode 100644 index 0000000..4c1e77e --- /dev/null +++ b/server/internal/nodes/store.go @@ -0,0 +1,158 @@ +package nodes + +import ( + "context" + "database/sql" + "fmt" + "time" + + agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1" +) + +// NodeRow holds a node's essential fields from the nodes table. +type NodeRow struct { + ID int64 + UUID string + Status string + RealityPBK string + RealitySNI string + Endpoint string + Hy2Port sql.NullInt32 +} + +// NodeStore is the persistence interface used by the nodes domain handlers. +// All methods are context-aware and safe for concurrent use. +type NodeStore interface { + // NodeByUUID looks up a node record by its UUID. + // Returns (nil, nil) when no matching row exists. + NodeByUUID(ctx context.Context, uuid string) (*NodeRow, error) + + // ConfigVersion returns the current global directory version. + // This is the version from the directory_version singleton table. + ConfigVersion(ctx context.Context) (int64, error) + + // ActiveNodeUUIDs returns the UUIDs of all nodes with status 'up' or 'draining'. + // Used by Broadcast to enumerate delivery targets. + ActiveNodeUUIDs(ctx context.Context) ([]string, error) + + // CredentialsForNode returns the active data-plane credentials for nodeUUID. + // Returns an empty slice until task 5d creates the connect_credentials table. + CredentialsForNode(ctx context.Context, nodeUUID string) ([]*agentv1.Credential, error) + + // UserIDByDpUUID maps a data-plane UUID to the owning user's internal ID. + // Returns (0, false, nil) if the dp_uuid is unknown or the user is inactive. + UserIDByDpUUID(ctx context.Context, dpUUID string) (int64, bool, error) + + // AccumulateUsage adds bytes and minutes to usage_daily for userID on date. + // Uses INSERT … ON DUPLICATE KEY UPDATE (idempotent within a day). + AccumulateUsage(ctx context.Context, userID int64, date time.Time, + bytesUp, bytesDown int64, minutes int64) error +} + +// SQLNodeStore implements NodeStore against a MySQL 8.x database. +type SQLNodeStore struct { + db *sql.DB +} + +// NewSQLNodeStore creates a SQLNodeStore backed by db. +func NewSQLNodeStore(db *sql.DB) *SQLNodeStore { + return &SQLNodeStore{db: db} +} + +// NodeByUUID looks up a node by UUID. Returns (nil, nil) when not found. +func (s *SQLNodeStore) NodeByUUID(ctx context.Context, uuid string) (*NodeRow, error) { + const q = ` + SELECT id, uuid, status, reality_pbk, reality_sni, endpoint, hy2_port + FROM nodes + WHERE uuid = ? + ` + var n NodeRow + err := s.db.QueryRowContext(ctx, q, uuid).Scan( + &n.ID, &n.UUID, &n.Status, + &n.RealityPBK, &n.RealitySNI, &n.Endpoint, &n.Hy2Port, + ) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("nodes.SQLNodeStore.NodeByUUID: %w", err) + } + return &n, nil +} + +// ConfigVersion returns the current global directory version. +// Returns 0 with nil error if the directory_version row does not yet exist. +func (s *SQLNodeStore) ConfigVersion(ctx context.Context) (int64, error) { + var v int64 + err := s.db.QueryRowContext(ctx, + `SELECT version FROM directory_version WHERE id = 1`).Scan(&v) + if err == sql.ErrNoRows { + return 0, nil + } + if err != nil { + return 0, fmt.Errorf("nodes.SQLNodeStore.ConfigVersion: %w", err) + } + return v, nil +} + +// ActiveNodeUUIDs returns UUIDs of all nodes with status 'up' or 'draining'. +func (s *SQLNodeStore) ActiveNodeUUIDs(ctx context.Context) ([]string, error) { + rows, err := s.db.QueryContext(ctx, + `SELECT uuid FROM nodes WHERE status IN ('up', 'draining')`) + if err != nil { + return nil, fmt.Errorf("nodes.SQLNodeStore.ActiveNodeUUIDs: %w", err) + } + defer rows.Close() + + var uuids []string + for rows.Next() { + var u string + if err := rows.Scan(&u); err != nil { + return nil, fmt.Errorf("nodes.SQLNodeStore.ActiveNodeUUIDs scan: %w", err) + } + uuids = append(uuids, u) + } + return uuids, rows.Err() +} + +// CredentialsForNode returns active credentials for nodeUUID. +// Stub: returns empty slice until task 5d creates the connect_credentials table. +func (s *SQLNodeStore) CredentialsForNode(_ context.Context, _ string) ([]*agentv1.Credential, error) { + return nil, nil +} + +// UserIDByDpUUID resolves a data-plane UUID to an active user's internal ID. +func (s *SQLNodeStore) UserIDByDpUUID(ctx context.Context, dpUUID string) (int64, bool, error) { + var userID int64 + err := s.db.QueryRowContext(ctx, + `SELECT id FROM users WHERE dp_uuid = ? AND status = 'active' LIMIT 1`, dpUUID, + ).Scan(&userID) + if err == sql.ErrNoRows { + return 0, false, nil + } + if err != nil { + return 0, false, fmt.Errorf("nodes.SQLNodeStore.UserIDByDpUUID: %w", err) + } + return userID, true, nil +} + +// AccumulateUsage adds bytes/minutes to usage_daily for the given user and date. +func (s *SQLNodeStore) AccumulateUsage( + ctx context.Context, userID int64, date time.Time, + bytesUp, bytesDown int64, minutes int64, +) error { + const q = ` + INSERT INTO usage_daily (user_id, date, bytes_up, bytes_down, minutes_used) + VALUES (?, ?, ?, ?, ?) + ON DUPLICATE KEY UPDATE + bytes_up = bytes_up + VALUES(bytes_up), + bytes_down = bytes_down + VALUES(bytes_down), + minutes_used = minutes_used + VALUES(minutes_used) + ` + if _, err := s.db.ExecContext(ctx, q, + userID, date.Format("2006-01-02"), bytesUp, bytesDown, minutes, + ); err != nil { + return fmt.Errorf("nodes.SQLNodeStore.AccumulateUsage: %w", err) + } + return nil +}