vdi-cloud/platform/internal/store/store.go

package store

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	_ "embed"
	"encoding/hex"
	"fmt"
	_ "github.com/jackc/pgx/v5/stdlib"
	"github.com/jmoiron/sqlx"
	"golang.org/x/crypto/bcrypt"
	"time"
	"vdi-cloud/platform/internal/model"
)

//go:embed schema.sql
var schema string

func Open(ctx context.Context, url string) (*sqlx.DB, error) {
	db, err := sqlx.ConnectContext(ctx, "pgx", url)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(12)
	db.SetMaxIdleConns(4)
	db.SetConnMaxLifetime(30 * time.Minute)
	if _, err = db.ExecContext(ctx, schema); err != nil {
		db.Close()
		return nil, err
	}
	return db, nil
}
func ID(prefix string) string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return prefix + hex.EncodeToString(b)
}
func Hash(secret string) string { h := sha256.Sum256([]byte(secret)); return hex.EncodeToString(h[:]) }
func Bootstrap(ctx context.Context, db *sqlx.DB, mode, username, password string) error {
	if _, err := db.ExecContext(ctx, "INSERT INTO platform_metadata(singleton,mode) VALUES(true,$1) ON CONFLICT DO NOTHING", mode); err != nil {
		return err
	}
	var existingMode string
	if err := db.GetContext(ctx, &existingMode, "SELECT mode FROM platform_metadata WHERE singleton=true"); err != nil {
		return err
	}
	if existingMode != mode {
		return fmt.Errorf("database mode is %s; use a separate database for %s", existingMode, mode)
	}
	admins := []struct{ username, name, role, password string }{{username, "平台管理员", "super_admin", password}}
	if mode == "demo" {
		admins = []struct{ username, name, role, password string }{{"admin", "系统管理员", "super_admin", "Demo@2026"}, {"operator", "运维人员", "operator", "Demo@2026"}, {"viewer", "只读人员", "viewer", "Demo@2026"}}
	}
	for _, a := range admins {
		var exists bool
		if err := db.GetContext(ctx, &exists, "SELECT EXISTS(SELECT 1 FROM administrators WHERE username=$1)", a.username); err != nil {
			return err
		}
		if exists {
			continue
		}
		if mode == "production" && len(a.password) < 12 {
			return fmt.Errorf("initial administrator password must contain at least 12 characters")
		}
		hash, err := bcrypt.GenerateFromPassword([]byte(a.password), bcrypt.DefaultCost)
		if err != nil {
			return err
		}
		if _, err = db.ExecContext(ctx, "INSERT INTO administrators(id,username,display_name,role,password_hash) VALUES($1,$2,$3,$4,$5) ON CONFLICT DO NOTHING", ID("adm_"), a.username, a.name, a.role, string(hash)); err != nil {
			return err
		}
	}
	if mode != "demo" {
		return nil
	}
	// Seed the demo only once; restarting must not restore revoked bindings.
	var seeded bool
	if err := db.GetContext(ctx, &seeded, "SELECT EXISTS(SELECT 1 FROM desktops WHERE provider='demo')"); err != nil {
		return err
	}
	if seeded {
		return nil
	}
	tx, err := db.BeginTxx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	statements := []string{
		`INSERT INTO desktops(id,name,os,node,vmid,pool,power_state,service_state,provider) VALUES ('desk_demo_windows','设计专属桌面','Windows','demo-node-a',201,'设计专属池','running','unknown','demo'),('desk_demo_ubuntu','研发办公桌面','Ubuntu','demo-node-a',202,'研发专属池','running','unknown','demo'),('desk_demo_kylin','国产办公桌面','麒麟','demo-node-a',203,'通用办公池','stopped','unknown','demo') ON CONFLICT DO NOTHING`,
		`INSERT INTO terminals(id,sn,name,group_name,status,agent_version,last_seen,bootstrap_hash) VALUES ('term_demo_001','DEMO-RK-001','设计工位 01','设计部','offline','演示版本',now()-interval '2 hours','disabled'),('term_demo_002','DEMO-RK-002','研发工位 02','研发部','offline','演示版本',now()-interval '3 hours','disabled') ON CONFLICT DO NOTHING`,
		`INSERT INTO users(id,username,display_name,terminal_id) VALUES ('user_demo_wang','wang.demo','王工','term_demo_001'),('user_demo_chen','chen.demo','陈工','term_demo_002') ON CONFLICT DO NOTHING`,
		`INSERT INTO bindings VALUES ('user_demo_wang','desk_demo_windows'),('user_demo_wang','desk_demo_ubuntu'),('user_demo_chen','desk_demo_ubuntu') ON CONFLICT DO NOTHING`,
		`INSERT INTO resources VALUES ('demo-node-a','demo-node-a','online',0.24,0.48,0.35,'demo',now()) ON CONFLICT DO NOTHING`,
	}
	for _, q := range statements {
		if _, err = tx.ExecContext(ctx, q); err != nil {
			return fmt.Errorf("demo seed: %w", err)
		}
	}
	return tx.Commit()
}

func Users(ctx context.Context, db *sqlx.DB) ([]model.User, error) {
	users := []model.User{}
	if err := db.SelectContext(ctx, &users, "SELECT * FROM users ORDER BY username"); err != nil {
		return nil, err
	}
	for i := range users {
		users[i].DesktopIDs = []string{}
		if err := db.SelectContext(ctx, &users[i].DesktopIDs, "SELECT desktop_id FROM bindings WHERE user_id=$1 ORDER BY desktop_id", users[i].ID); err != nil {
			return nil, err
		}
	}
	return users, nil
}