vdi-cloud/platform/internal/server/catalog.go

package server

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"regexp"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/jmoiron/sqlx"
	"vdi-cloud/platform/internal/model"
	"vdi-cloud/platform/internal/store"
	"vdi-cloud/platform/pkg/proxmox"
)

func (s *Server) discoverTemplates(c *gin.Context) {
	pve := s.currentPVE()
	if pve == nil {
		s.fail(c, 503, "pve_unconfigured", "尚未配置 Proxmox")
		return
	}
	resources, err := pve.Resources(c.Request.Context())
	if err != nil {
		s.fail(c, 502, "pve_unavailable", "无法发现模板，请检查 PVE 接入")
		return
	}
	items := []gin.H{}
	for _, r := range resources {
		if r.Type == "qemu" && r.Template == 1 {
			items = append(items, gin.H{"name": r.Name, "node": r.Node, "vmid": r.VMID, "status": r.Status})
		}
	}
	list(c, items)
}
func (s *Server) templateList(c *gin.Context) {
	items := []model.ImageTemplate{}
	if err := s.DB.SelectContext(c.Request.Context(), &items, "SELECT * FROM image_templates ORDER BY name,version"); err != nil {
		s.dbError(c, err)
		return
	}
	list(c, items)
}
func (s *Server) registerTemplate(c *gin.Context) {
	pve := s.currentPVE()
	var in struct {
		Name    string `json:"name" binding:"required,max=80"`
		Version string `json:"version" binding:"required,max=40"`
		OS      string `json:"os" binding:"required,oneof=Windows Ubuntu 麒麟"`
		Node    string `json:"node" binding:"required,max=100"`
		VMID    int    `json:"vmid" binding:"min=100,max=999999999"`
	}
	if !s.decode(c, &in) {
		return
	}
	in.Name = strings.TrimSpace(in.Name)
	in.Version = strings.TrimSpace(in.Version)
	if in.Name == "" || in.Version == "" {
		s.fail(c, 400, "invalid_request", "模板名称和版本不能为空")
		return
	}
	if pve == nil {
		s.fail(c, 503, "pve_unconfigured", "模板登记需要先配置 Proxmox")
		return
	}
	config, err := pve.VMConfig(c.Request.Context(), in.Node, in.VMID)
	if err != nil {
		s.fail(c, 502, "pve_unavailable", "无法读取模板配置，请核对节点、VMID 和 VM.Audit 权限")
		return
	}
	if _, err = proxmox.TemplateDisks(config); err != nil {
		s.fail(c, 400, "invalid_template", err.Error())
		return
	}
	digest, _ := config["digest"].(string)
	if digest == "" {
		s.fail(c, 502, "invalid_template", "PVE 模板配置缺少版本摘要")
		return
	}
	var item model.ImageTemplate
	err = s.DB.GetContext(c.Request.Context(), &item, `INSERT INTO image_templates(id,name,version,os,node,vmid,digest) VALUES($1,$2,$3,$4,$5,$6,$7) RETURNING *`, store.ID("tpl_"), in.Name, in.Version, in.OS, in.Node, in.VMID, digest)
	if conflict(err) {
		s.fail(c, 409, "template_conflict", "该模板或名称版本已经登记")
		return
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	s.log(c, "register_template", item.ID, "success")
	c.JSON(201, item)
}
func (s *Server) updateTemplate(c *gin.Context) {
	var in struct {
		Enabled *bool `json:"enabled" binding:"required"`
	}
	if !s.decode(c, &in) {
		return
	}
	var item model.ImageTemplate
	if err := s.DB.GetContext(c.Request.Context(), &item, "UPDATE image_templates SET enabled=$1 WHERE id=$2 RETURNING *", *in.Enabled, c.Param("id")); err != nil {
		s.dbError(c, err)
		return
	}
	s.log(c, "update_template", item.ID, "success")
	c.JSON(200, item)
}

const poolSelect = `SELECT p.*,
 (SELECT count(*) FROM desktops d WHERE d.pool_id=p.id) desktop_count,
 (SELECT count(*) FROM provisioning_jobs j WHERE j.pool_id=p.id AND j.state IN ('reserved','uncertain')) reserved_count,
 (SELECT count(*) FROM desktops d WHERE d.pool_id=p.id AND EXISTS(SELECT 1 FROM bindings b JOIN users u ON u.id=b.user_id WHERE b.desktop_id=d.id AND u.enabled)) assigned_count,
 (SELECT count(*) FROM desktops d WHERE d.pool_id=p.id AND EXISTS(SELECT 1 FROM terminals t WHERE t.session_desktop_id=d.id AND t.last_seen>now()-interval '90 seconds')) active_count
 FROM desktop_pools p`

func (s *Server) poolList(c *gin.Context) {
	items := []model.DesktopPool{}
	if err := s.DB.SelectContext(c.Request.Context(), &items, poolSelect+" ORDER BY p.name"); err != nil {
		s.dbError(c, err)
		return
	}
	list(c, items)
}

type poolInput struct {
	Name       string  `json:"name" binding:"required,max=80"`
	Type       string  `json:"type" binding:"required,oneof=dedicated shared"`
	MaxCount   int     `json:"maxCount" binding:"min=1,max=10000"`
	TemplateID *string `json:"templateId"`
	Enabled    *bool   `json:"enabled" binding:"required"`
}

// The catalog lock serializes capacity reservations, pool edits and membership changes.
func catalogLock(ctx context.Context, tx *sqlx.Tx) error {
	_, err := tx.ExecContext(ctx, "SELECT pg_advisory_xact_lock(7651092026)")
	return err
}
func (s *Server) savePool(c *gin.Context) {
	var in poolInput
	if !s.decode(c, &in) {
		return
	}
	in.Name = strings.TrimSpace(in.Name)
	if in.Name == "" {
		s.fail(c, 400, "invalid_request", "桌面池名称不能为空")
		return
	}
	if in.TemplateID != nil && *in.TemplateID == "" {
		in.TemplateID = nil
	}
	ctx := c.Request.Context()
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		s.dbError(c, err)
		return
	}
	defer tx.Rollback()
	if err = catalogLock(ctx, tx); err != nil {
		s.dbError(c, err)
		return
	}
	id := c.Param("id")
	if id == "" {
		id = store.ID("pool_")
		_, err = tx.ExecContext(ctx, "INSERT INTO desktop_pools(id,name,type,max_count,template_id,enabled) VALUES($1,$2,$3,$4,$5,$6)", id, in.Name, in.Type, in.MaxCount, in.TemplateID, *in.Enabled)
	} else {
		var current model.DesktopPool
		if err = tx.GetContext(ctx, &current, poolSelect+" WHERE p.id=$1", id); err != nil {
			s.dbError(c, err)
			return
		}
		if in.MaxCount < current.DesktopCount+current.ReservedCount {
			s.fail(c, 409, "pool_capacity", "容量不能小于现有桌面与待确认创建任务的总数")
			return
		}
		if in.Type != current.Type && current.DesktopCount+current.ReservedCount > 0 {
			s.fail(c, 409, "pool_in_use", "非空桌面池不能切换公共/专属类型")
			return
		}
		_, err = tx.ExecContext(ctx, "UPDATE desktop_pools SET name=$1,type=$2,max_count=$3,template_id=$4,enabled=$5 WHERE id=$6", in.Name, in.Type, in.MaxCount, in.TemplateID, *in.Enabled, id)
		if err == nil {
			_, err = tx.ExecContext(ctx, "UPDATE desktops SET pool=$1 WHERE pool_id=$2", in.Name, id)
		}
	}
	if conflict(err) {
		s.fail(c, 409, "pool_conflict", "名称已存在或所选模板不存在")
		return
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	var item model.DesktopPool
	if err = tx.GetContext(ctx, &item, poolSelect+" WHERE p.id=$1", id); err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	s.log(c, "save_pool", id, "success")
	c.JSON(200, item)
}
func (s *Server) poolMembers(c *gin.Context) {
	var in struct {
		DesktopIDs []string `json:"desktopIds" binding:"max=1000"`
	}
	if !s.decode(c, &in) {
		return
	}
	ctx := c.Request.Context()
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		s.dbError(c, err)
		return
	}
	defer tx.Rollback()
	if err = catalogLock(ctx, tx); err != nil {
		s.dbError(c, err)
		return
	}
	var pool model.DesktopPool
	if err = tx.GetContext(ctx, &pool, poolSelect+" WHERE p.id=$1", c.Param("id")); err != nil {
		s.dbError(c, err)
		return
	}
	ids := map[string]bool{}
	for _, id := range in.DesktopIDs {
		ids[id] = true
	}
	if len(ids)+pool.ReservedCount > pool.MaxCount {
		s.fail(c, 409, "pool_capacity", "桌面数量与创建预留超过池容量")
		return
	}
	for id := range ids {
		var desktop model.Desktop
		if err = tx.GetContext(ctx, &desktop, "SELECT * FROM desktops WHERE id=$1", id); err != nil {
			s.fail(c, 400, "invalid_desktop", "所选桌面不存在")
			return
		}
		if desktop.PoolID != nil && *desktop.PoolID != pool.ID {
			s.fail(c, 409, "pool_membership", "桌面已属于其他池，请先从原池移出")
			return
		}
	}
	if _, err = tx.ExecContext(ctx, "UPDATE desktops SET pool_id=NULL,pool='' WHERE pool_id=$1", pool.ID); err != nil {
		s.dbError(c, err)
		return
	}
	for id := range ids {
		if _, err = tx.ExecContext(ctx, "UPDATE desktops SET pool_id=$1,pool=$2 WHERE id=$3", pool.ID, pool.Name, id); err != nil {
			s.dbError(c, err)
			return
		}
	}
	if err = tx.Commit(); err != nil {
		s.dbError(c, err)
		return
	}
	s.log(c, "update_pool_members", pool.ID, "success")
	c.JSON(200, gin.H{"accepted": true})
}

type cloneInput struct {
	TemplateID string  `json:"templateId" binding:"required,max=100"`
	PoolID     *string `json:"poolId"`
	Name       string  `json:"name" binding:"required,max=80"`
	VMID       int     `json:"vmid" binding:"min=100,max=999999999"`
	Full       bool    `json:"full"`
	Storage    string  `json:"storage" binding:"max=100"`
}
type provisioningJob struct {
	TaskID      string    `db:"task_id"`
	TemplateID  string    `db:"template_id"`
	PoolID      *string   `db:"pool_id"`
	DesktopID   string    `db:"desktop_id"`
	VMID        int       `db:"vmid"`
	RequestJSON string    `db:"request_json"`
	RequestHash string    `db:"request_hash"`
	State       string    `db:"state"`
	CreatedAt   time.Time `db:"created_at"`
}

func normalizeClone(in *cloneInput) {
	in.Name = strings.TrimSpace(in.Name)
	if in.PoolID != nil && *in.PoolID == "" {
		in.PoolID = nil
	}
}
func (s *Server) clonePreflight(ctx context.Context, in cloneInput) (model.ImageTemplate, error) {
	pve := s.currentPVE()
	var template model.ImageTemplate
	if pve == nil {
		return template, errors.New("尚未配置 Proxmox")
	}
	if in.Name == "" || in.Full && in.Storage == "" || !in.Full && in.Storage != "" {
		return template, errors.New("请填写桌面名称；完整克隆必须选择存储，链接克隆沿用模板存储")
	}
	if err := s.DB.GetContext(ctx, &template, "SELECT * FROM image_templates WHERE id=$1 AND enabled", in.TemplateID); err != nil {
		return template, errors.New("模板不存在或已停用")
	}
	config, err := pve.VMConfig(ctx, template.Node, template.VMID)
	if err != nil {
		return template, errors.New("无法读取源模板配置")
	}
	disks, err := proxmox.TemplateDisks(config)
	if err != nil {
		return template, err
	}
	if config["digest"] != template.Digest {
		return template, errors.New("源模板已变更，请重新验证模板后再创建")
	}
	can, err := pve.CanClone(ctx, template.Node, template.VMID, in.Full)
	if err != nil || !can {
		return template, errors.New("模板所在存储不支持所选克隆方式")
	}
	resources, err := pve.Resources(ctx)
	if err != nil {
		return template, errors.New("无法检查目标 VMID")
	}
	for _, r := range resources {
		if (r.Type == "qemu" || r.Type == "lxc") && r.VMID == in.VMID {
			return template, errors.New("目标 VMID 已被 PVE 使用")
		}
	}
	for _, permission := range []struct{ path, name string }{{fmt.Sprintf("/vms/%d", template.VMID), "VM.Clone"}, {fmt.Sprintf("/vms/%d", in.VMID), "VM.Allocate"}} {
		ok, err := pve.HasPermission(ctx, permission.path, permission.name)
		if err != nil || !ok {
			return template, fmt.Errorf("PVE 平台账号缺少 %s 权限（%s）", permission.name, permission.path)
		}
	}
	storages, err := pve.Storages(ctx, template.Node)
	if err != nil {
		return template, errors.New("无法查询目标存储")
	}
	if in.Full {
		total := float64(0)
		for _, size := range disks {
			total += size
		}
		disks = map[string]float64{in.Storage: total}
	}
	for name, size := range disks {
		found := false
		for _, storage := range storages {
			if storage.Name == name && storage.Active == 1 && storage.Enabled == 1 && strings.Contains(","+storage.Content+",", ",images,") {
				found = true
				needed := float64(1 << 30)
				if in.Full {
					needed += size
				}
				if storage.Available < needed {
					return template, fmt.Errorf("存储 %s 可用空间不足，已预留 1 GiB 余量", name)
				}
			}
		}
		if !found {
			return template, fmt.Errorf("存储 %s 不可用或不支持虚拟机磁盘", name)
		}
		ok, err := pve.HasPermission(ctx, "/storage/"+name, "Datastore.AllocateSpace")
		if err != nil || !ok {
			return template, fmt.Errorf("存储 %s 缺少 Datastore.AllocateSpace 权限", name)
		}
	}
	return template, nil
}
func (s *Server) storageList(c *gin.Context) {
	pve := s.currentPVE()
	if pve == nil {
		s.fail(c, 503, "pve_unconfigured", "尚未配置 Proxmox")
		return
	}
	node := c.Query("node")
	if !regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9.-]{0,99}$`).MatchString(node) {
		s.fail(c, 400, "invalid_node", "请选择节点")
		return
	}
	items, err := pve.Storages(c.Request.Context(), node)
	if err != nil {
		s.fail(c, 502, "pve_unavailable", "存储查询失败")
		return
	}
	list(c, items)
}
func (s *Server) checkClone(c *gin.Context) {
	var in cloneInput
	if !s.decode(c, &in) {
		return
	}
	normalizeClone(&in)
	_, err := s.clonePreflight(c.Request.Context(), in)
	if err != nil {
		c.JSON(200, gin.H{"available": false, "message": err.Error()})
		return
	}
	if in.PoolID != nil {
		var pool model.DesktopPool
		if err = s.DB.GetContext(c.Request.Context(), &pool, poolSelect+" WHERE p.id=$1", *in.PoolID); err != nil {
			c.JSON(200, gin.H{"available": false, "message": "所选桌面池不存在"})
			return
		}
		if !pool.Enabled || pool.DesktopCount+pool.ReservedCount >= pool.MaxCount || (pool.TemplateID != nil && *pool.TemplateID != in.TemplateID) {
			c.JSON(200, gin.H{"available": false, "message": "桌面池已停用、容量不足或模板不匹配"})
			return
		}
	}
	var reserved bool
	if err = s.DB.GetContext(c.Request.Context(), &reserved, "SELECT EXISTS(SELECT 1 FROM provisioning_jobs WHERE vmid=$1)", in.VMID); err != nil {
		s.dbError(c, err)
		return
	}
	if reserved {
		c.JSON(200, gin.H{"available": false, "message": "目标 VMID 已有创建记录，请核对任务并选择其他 VMID"})
		return
	}
	c.JSON(200, gin.H{"available": true, "message": "计算侧预检通过；提交时会再次检查容量和模板，创建后需要完成系统初始化与配对"})
}
func (s *Server) cloneDesktop(c *gin.Context) {
	var in cloneInput
	if !s.decode(c, &in) {
		return
	}
	normalizeClone(&in)
	key := c.GetHeader("Idempotency-Key")
	if key == "" || len(key) > 100 {
		s.fail(c, 400, "idempotency_required", "请提供 Idempotency-Key")
		return
	}
	ctx := c.Request.Context()
	body, _ := json.Marshal(in)
	hash := store.Hash(string(body))
	actor := s.admin(c).Username
	var existing model.Task
	if err := s.DB.GetContext(ctx, &existing, "SELECT * FROM tasks WHERE actor=$1 AND idempotency_key=$2", actor, key); err == nil {
		var same bool
		_ = s.DB.GetContext(ctx, &same, "SELECT EXISTS(SELECT 1 FROM provisioning_jobs WHERE task_id=$1 AND request_hash=$2)", existing.ID, hash)
		if !same {
			s.fail(c, 409, "idempotency_conflict", "相同请求标识不能用于不同创建参数")
			return
		}
		c.JSON(202, existing)
		return
	} else if !errors.Is(err, sql.ErrNoRows) {
		s.dbError(c, err)
		return
	}
	_, err := s.clonePreflight(ctx, in)
	if err != nil {
		s.fail(c, 409, "clone_preflight_failed", err.Error())
		return
	}
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		s.dbError(c, err)
		return
	}
	defer tx.Rollback()
	if err = catalogLock(ctx, tx); err != nil {
		s.dbError(c, err)
		return
	}
	var enabled bool
	if err = tx.GetContext(ctx, &enabled, "SELECT enabled FROM image_templates WHERE id=$1 FOR SHARE", in.TemplateID); err != nil || !enabled {
		s.fail(c, 409, "template_disabled", "模板已停用")
		return
	}
	if in.PoolID != nil {
		var pool model.DesktopPool
		if err = tx.GetContext(ctx, &pool, poolSelect+" WHERE p.id=$1", *in.PoolID); err != nil {
			s.fail(c, 400, "invalid_pool", "所选桌面池不存在")
			return
		}
		if !pool.Enabled || pool.DesktopCount+pool.ReservedCount >= pool.MaxCount {
			s.fail(c, 409, "pool_capacity", "桌面池已停用或没有剩余容量")
			return
		}
		if pool.TemplateID != nil && *pool.TemplateID != in.TemplateID {
			s.fail(c, 409, "pool_template", "所选模板与池的模板不一致")
			return
		}
	}
	var taken bool
	if err = tx.GetContext(ctx, &taken, "SELECT EXISTS(SELECT 1 FROM desktops WHERE provider='proxmox' AND vmid=$1)", in.VMID); err != nil {
		s.dbError(c, err)
		return
	}
	if taken {
		s.fail(c, 409, "vmid_conflict", "目标 VMID 已纳管")
		return
	}
	var task model.Task
	desktopID := store.ID("desk_")
	err = tx.GetContext(ctx, &task, `INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,message,idempotency_key) VALUES($1,'provision','clone',$2,$3,$4,'queued','创建请求已排队，等待 PVE 最终结果',$5) RETURNING *`, store.ID("task_"), desktopID, in.Name, actor, key)
	if err == nil {
		_, err = tx.ExecContext(ctx, `INSERT INTO provisioning_jobs(task_id,template_id,pool_id,desktop_id,vmid,request_json,request_hash) VALUES($1,$2,$3,$4,$5,$6,$7)`, task.ID, in.TemplateID, in.PoolID, desktopID, in.VMID, string(body), hash)
	}
	if conflict(err) {
		s.fail(c, 409, "provision_conflict", "该 VMID 已有创建记录或重复请求，请核对任务")
		return
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	_, err = tx.ExecContext(ctx, "INSERT INTO audit_logs VALUES($1,$2,'clone',$3,'submitted',$4,now())", store.ID("audit_"), actor, in.Name, c.GetString("requestId"))
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(202, task)
}

func validEndpoint(endpoint string) bool {
	if net.ParseIP(endpoint) != nil {
		return true
	}
	if len(endpoint) > 253 {
		return false
	}
	for _, label := range strings.Split(endpoint, ".") {
		if !regexp.MustCompile(`^[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?$`).MatchString(label) {
			return false
		}
	}
	return true
}
func (s *Server) configureDesktop(c *gin.Context) {
	var in struct {
		Endpoint    string `json:"endpoint" binding:"required,max=253"`
		Width       int    `json:"width" binding:"min=640,max=7680"`
		Height      int    `json:"height" binding:"min=480,max=4320"`
		FPS         int    `json:"fps" binding:"min=10,max=120"`
		Initialized bool   `json:"initialized"`
	}
	if !s.decode(c, &in) {
		return
	}
	if !validEndpoint(in.Endpoint) {
		s.fail(c, 400, "invalid_endpoint", "请填写有效的 IP 或主机名")
		return
	}
	ctx := c.Request.Context()
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		s.dbError(c, err)
		return
	}
	defer tx.Rollback()
	if !s.lockAuthorization(c, tx) {
		return
	}
	var desktop model.Desktop
	if err = tx.GetContext(ctx, &desktop, "SELECT * FROM desktops WHERE id=$1 FOR UPDATE", c.Param("id")); err != nil {
		s.dbError(c, err)
		return
	}
	var active bool
	if err = tx.GetContext(ctx, &active, "SELECT EXISTS(SELECT 1 FROM terminals WHERE session_desktop_id=$1 AND last_seen>now()-interval '90 seconds')", desktop.ID); err != nil {
		s.dbError(c, err)
		return
	}
	if active {
		s.fail(c, 409, "desktop_busy", "桌面有连接中或活动会话，请先结束连接")
		return
	}
	if blocked, e := maintenanceBlocked(ctx, tx, desktop.ID, "", true); e != nil {
		s.dbError(c, e)
		return
	} else if blocked {
		s.fail(c, 409, "desktop_maintenance", "桌面维护任务进行中或结果未确认，请先核对")
		return
	}
	if desktop.SetupRequired && !in.Initialized {
		s.fail(c, 400, "initialization_required", "请先确认系统身份、IP、Sunshine 服务和终端配对已经完成")
		return
	}
	err = tx.GetContext(ctx, &desktop, `UPDATE desktops SET endpoint=$1,width=$2,height=$3,fps=$4,setup_required=false,service_state='unknown',updated_at=now() WHERE id=$5 RETURNING *`, in.Endpoint, in.Width, in.Height, in.FPS, desktop.ID)
	if err == nil {
		_, err = tx.ExecContext(ctx, "DELETE FROM connection_tickets WHERE desktop_id=$1", desktop.ID)
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	s.log(c, "configure_desktop", desktop.ID, "success")
	c.JSON(200, desktop)
}