feat(nodes): gRPC server + Hub + Redis pub/sub cross-instance delivery (tsk__58l3wTLvaSn)
Implements AgentService gRPC server with all 6 RPCs (Enroll/Register/Heartbeat/ Subscribe/Ack/ReportUsage), Hub command routing with Redis ZSET at-least-once persistence and cross-instance pub/sub delivery, LoadCache for node:load metrics, NodeStore SQL interface + MySQL implementation, and full mTLS gRPC listener in main. Integration tests: 18 tests covering full Enroll→Register→Heartbeat→Subscribe→Ack flow, reconnect resume with last_command_id, and cross-instance pub/sub delivery via bufconn + miniredis + mockNodeStore + real mTLS certificates. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
wangjia
+103
-9
@@ -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
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user