pangolin/server/internal/nodes/handler_grpc.go

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.
	// reality_prk is the PRIVATE key the agent's VLESS inbound needs;
	// reality_pbk is the PUBLIC key sent to clients in the connect config.
	if node.RealityPRK != "" {
		snap.Reality = &agentv1.RealityInbound{
			PrivateKey: node.RealityPRK,
			ShortID:    node.RealityShortID,
			ServerName: node.RealitySNI,
		}
	} else if node.RealityPBK != "" {
		// Fallback for nodes seeded before migration 000011: use pbk field.
		snap.Reality = &agentv1.RealityInbound{
			PrivateKey: node.RealityPBK,
			ShortID:    node.RealityShortID,
			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
}