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 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 }