pangolin/server/internal/nodes/store.go

package nodes

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	dbx "github.com/wangjia/pangolin/server/internal/db"
	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
	Region       string
	NameZH       string
	NameEN       string
	Tier         string
	Status       string
	RealityPBK   string // REALITY x25519 PUBLIC key (for client connect config)
	RealityPRK   string // REALITY x25519 PRIVATE key (for agent inbound TLS)
	RealitySNI   string
	RealityShortID string
	Endpoint     string
	Hy2Port      sql.NullInt32
}

// Entitlement summarises a user's active plan for the connect gate.
type Entitlement struct {
	DpUUID       string
	PlanCode     string
	AdGate       bool // true = free plan, require ad unlock + minute quota
	DailyMinutes sql.NullInt64
	ExpiresAt    sql.NullTime // latest subscription expiry (nil = trial/active)
}

// 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)

	// ListUp returns all nodes with status='up', ordered by weight DESC.
	ListUp(ctx context.Context) ([]*NodeRow, error)

	// EntitlementForUser returns the user's dp_uuid and active plan entitlement.
	// Returns (nil, nil) when the user has no active subscription (treats as free).
	EntitlementForUser(ctx context.Context, userID int64) (*Entitlement, 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.
	CredentialsForNode(ctx context.Context, nodeUUID string) ([]*agentv1.Credential, error)

	// PersistCredential upserts a credential row in connect_credentials.
	PersistCredential(ctx context.Context, nodeID int64, cred *agentv1.Credential, expiresAt time.Time) error

	// DeleteCredential removes the credential for (nodeID, dpUUID).
	DeleteCredential(ctx context.Context, nodeID int64, dpUUID string) 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 SQL database (MySQL or SQLite).
type SQLNodeStore struct {
	db      *sql.DB
	dialect dbx.Dialect
}

// NewSQLNodeStore creates a SQLNodeStore backed by db.
func NewSQLNodeStore(db *sql.DB) *SQLNodeStore {
	return &SQLNodeStore{db: db, dialect: dbx.DialectForDB(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, region, name_zh, name_en, tier, status,
		       reality_pbk, reality_prk, reality_sni, reality_short_id,
		       endpoint, hy2_port
		FROM nodes
		WHERE uuid = ?
	`
	var n NodeRow
	err := s.db.QueryRowContext(ctx, q, uuid).Scan(
		&n.ID, &n.UUID, &n.Region, &n.NameZH, &n.NameEN, &n.Tier, &n.Status,
		&n.RealityPBK, &n.RealityPRK, &n.RealitySNI, &n.RealityShortID,
		&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
}

// ListUp returns nodes with status='up', ordered by weight DESC.
func (s *SQLNodeStore) ListUp(ctx context.Context) ([]*NodeRow, error) {
	const q = `
		SELECT id, uuid, region, name_zh, name_en, tier, status,
		       reality_pbk, reality_prk, reality_sni, reality_short_id,
		       endpoint, hy2_port
		FROM nodes
		WHERE status = 'up'
		ORDER BY weight DESC
	`
	rows, err := s.db.QueryContext(ctx, q)
	if err != nil {
		return nil, fmt.Errorf("nodes.SQLNodeStore.ListUp: %w", err)
	}
	defer rows.Close()

	var out []*NodeRow
	for rows.Next() {
		var n NodeRow
		if err := rows.Scan(
			&n.ID, &n.UUID, &n.Region, &n.NameZH, &n.NameEN, &n.Tier, &n.Status,
			&n.RealityPBK, &n.RealityPRK, &n.RealitySNI, &n.RealityShortID,
			&n.Endpoint, &n.Hy2Port,
		); err != nil {
			return nil, fmt.Errorf("nodes.SQLNodeStore.ListUp scan: %w", err)
		}
		out = append(out, &n)
	}
	return out, rows.Err()
}

// EntitlementForUser returns the user's dp_uuid and best active plan entitlement.
// Picks the subscription with the latest expires_at; falls back to the 'free' plan
// if the user has no active subscription.
func (s *SQLNodeStore) EntitlementForUser(ctx context.Context, userID int64) (*Entitlement, error) {
	// First get dp_uuid from the users table.
	var dpUUID string
	if err := s.db.QueryRowContext(ctx,
		`SELECT dp_uuid FROM users WHERE id = ? AND status = 'active'`, userID,
	).Scan(&dpUUID); err == sql.ErrNoRows {
		return nil, nil
	} else if err != nil {
		return nil, fmt.Errorf("nodes.SQLNodeStore.EntitlementForUser: dp_uuid: %w", err)
	}

	// Look up the best active subscription.
	const q = `
		SELECT p.code, p.ad_gate, p.daily_minutes, s.expires_at
		FROM subscriptions s
		JOIN plans p ON p.id = s.plan_id
		WHERE s.user_id = ? AND s.expires_at > ?
		ORDER BY s.expires_at DESC
		LIMIT 1
	`
	e := &Entitlement{DpUUID: dpUUID}
	err := s.db.QueryRowContext(ctx, q, userID, time.Now().UTC()).Scan(
		&e.PlanCode, &e.AdGate, &e.DailyMinutes, &e.ExpiresAt,
	)
	if err == sql.ErrNoRows {
		// No active subscription → free plan defaults.
		e.PlanCode = "free"
		e.AdGate = true
		e.DailyMinutes = sql.NullInt64{Valid: true, Int64: 10}
		return e, nil
	}
	if err != nil {
		return nil, fmt.Errorf("nodes.SQLNodeStore.EntitlementForUser: plan: %w", err)
	}
	return e, 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 (non-expired) credentials for nodeUUID.
func (s *SQLNodeStore) CredentialsForNode(ctx context.Context, nodeUUID string) ([]*agentv1.Credential, error) {
	const q = `
		SELECT cc.dp_uuid, cc.protocol, cc.flow, cc.expires_at
		FROM connect_credentials cc
		JOIN nodes n ON n.id = cc.node_id
		WHERE n.uuid = ? AND cc.expires_at > ?
	`
	rows, err := s.db.QueryContext(ctx, q, nodeUUID, time.Now().UTC())
	if err != nil {
		return nil, fmt.Errorf("nodes.SQLNodeStore.CredentialsForNode: %w", err)
	}
	defer rows.Close()

	var out []*agentv1.Credential
	for rows.Next() {
		var c agentv1.Credential
		var expiresAt time.Time
		if err := rows.Scan(&c.DpUUID, &c.Protocol, &c.Flow, &expiresAt); err != nil {
			return nil, fmt.Errorf("nodes.SQLNodeStore.CredentialsForNode scan: %w", err)
		}
		c.ExpiresAtUnix = expiresAt.Unix()
		out = append(out, &c)
	}
	return out, rows.Err()
}

// PersistCredential upserts the credential into connect_credentials.
func (s *SQLNodeStore) PersistCredential(ctx context.Context, nodeID int64, cred *agentv1.Credential, expiresAt time.Time) error {
	q := `
		INSERT INTO connect_credentials (node_id, dp_uuid, protocol, flow, expires_at)
		VALUES (?, ?, ?, ?, ?)
		` + s.dialect.Upsert([]string{"node_id", "dp_uuid"},
		"protocol = EXCLUDED.protocol",
		"flow = EXCLUDED.flow",
		"expires_at = EXCLUDED.expires_at")
	if _, err := s.db.ExecContext(ctx, q,
		nodeID, cred.DpUUID, int32(cred.Protocol), cred.Flow, expiresAt.UTC(),
	); err != nil {
		return fmt.Errorf("nodes.SQLNodeStore.PersistCredential: %w", err)
	}
	return nil
}

// DeleteCredential removes the credential for (nodeID, dpUUID).
func (s *SQLNodeStore) DeleteCredential(ctx context.Context, nodeID int64, dpUUID string) error {
	if _, err := s.db.ExecContext(ctx,
		`DELETE FROM connect_credentials WHERE node_id = ? AND dp_uuid = ?`,
		nodeID, dpUUID,
	); err != nil {
		return fmt.Errorf("nodes.SQLNodeStore.DeleteCredential: %w", err)
	}
	return 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 {
	q := `
		INSERT INTO usage_daily (user_id, date, bytes_up, bytes_down, minutes_used)
		VALUES (?, ?, ?, ?, ?)
		` + s.dialect.Upsert([]string{"user_id", "date"},
		"bytes_up = bytes_up + EXCLUDED.bytes_up",
		"bytes_down = bytes_down + EXCLUDED.bytes_down",
		"minutes_used = minutes_used + EXCLUDED.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
}