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

package server

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"sort"
	"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"
)

type snapshotInput struct {
	Action      string `json:"action" binding:"required,oneof=create delete rollback"`
	Name        string `json:"name" binding:"required,max=40"`
	Description string `json:"description" binding:"max=500"`
	Digest      string `json:"digest" binding:"required,max=80"`
	Revision    string `json:"revision" binding:"max=80"`
	ConfirmVMID int    `json:"confirmVmid"`
}
type snapshotItem struct {
	Name           string `json:"name"`
	Description    string `json:"description"`
	CreatedAt      int64  `json:"createdAt"`
	Parent         string `json:"parent"`
	IncludesMemory bool   `json:"includesMemory"`
	State          string `json:"state"`
	Revision       string `json:"revision"`
}

func snapshotRevision(item proxmox.Snapshot) string {
	raw, _ := json.Marshal(item)
	return store.Hash(string(raw))
}
func snapshotConfigHash(config map[string]any) string {
	// PVE refreshes vmgenid on rollback. Other fields here describe the snapshot
	// tree or transient execution state, rather than the guest configuration.
	skip := map[string]bool{"digest": true, "parent": true, "snaptime": true, "snapstate": true, "vmstate": true, "description": true, "snapshots": true, "pending": true, "lock": true, "running": true, "vmgenid": true, "runningmachine": true, "runningcpu": true, "running-nets-host-mtu": true}
	out := map[string]any{}
	for k, v := range config {
		if !skip[k] {
			out[k] = fmt.Sprint(v)
		}
	}
	raw, _ := json.Marshal(out)
	return store.Hash(string(raw))
}
func snapshotBaseBlockers(h proxmox.VMHardware, config map[string]any) ([]string, []string) {
	blockers := append([]string{}, h.Blockers...)
	if h.PowerState != "stopped" {
		blockers = append(blockers, "请先正常关闭桌面，再管理快照")
	}
	disks, e := proxmox.SnapshotDisks(config)
	scope := []string{}
	if e != nil {
		blockers = append(blockers, e.Error())
	} else {
		for key := range disks {
			scope = append(scope, key)
		}
		sort.Strings(scope)
	}
	return blockers, scope
}
func (s *Server) snapshotView(c *gin.Context) {
	ctx := c.Request.Context()
	d, h, config, err := s.loadHardware(ctx, c.Param("id"), s.DB)
	if err != nil {
		s.fail(c, 502, "snapshots_unavailable", "无法读取真实桌面配置，请检查 PVE 接入和 VM.Audit 权限")
		return
	}
	items, err := s.currentPVE().Snapshots(ctx, d.Node, d.VMID)
	if err != nil {
		s.fail(c, 502, "snapshots_unavailable", "无法读取 PVE 快照，请检查 VM.Audit 权限")
		return
	}
	blockers, scope := snapshotBaseBlockers(h, config)
	blocked, err := maintenanceBlocked(ctx, s.DB, d.ID, "", true)
	if err != nil {
		s.dbError(c, err)
		return
	}
	if blocked {
		blockers = append(blockers, "有桌面维护任务正在执行或结果未确认，请先核对任务")
	}
	permission := func(name string) bool {
		ok, e := s.currentPVE().HasPermission(ctx, fmt.Sprintf("/vms/%d", d.VMID), name)
		return e == nil && ok
	}
	manage, rollback := permission("VM.Snapshot"), permission("VM.Snapshot.Rollback")
	out := []snapshotItem{}
	for _, item := range items {
		out = append(out, snapshotItem{item.Name, item.Description, item.Time, item.Parent, item.Memory != 0, item.State, snapshotRevision(item)})
	}
	c.JSON(200, gin.H{"items": out, "digest": h.Digest, "powerState": h.PowerState, "blockers": blockers, "disks": scope, "canCreate": manage, "canDelete": manage, "canRestore": manage || rollback})
}

var snapshotName = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_-]{0,39}$`)

func (s *Server) snapshotPreflight(ctx context.Context, db sqlx.ExtContext, id string, in snapshotInput, own string) (model.Desktop, string, error) {
	var d model.Desktop
	if !snapshotName.MatchString(in.Name) || strings.EqualFold(in.Name, "current") || strings.EqualFold(in.Name, "pending") {
		return d, "", errors.New("快照名称使用字母开头的 1～40 位字母、数字、下划线或短横线，不能使用保留名称")
	}
	if in.Action != "create" && in.Description != "" || in.Action != "rollback" && in.ConfirmVMID != 0 {
		return d, "", errors.New("请求包含不适用于当前操作的参数")
	}
	d, h, config, err := s.loadHardware(ctx, id, db)
	if err != nil {
		return d, "", errors.New("无法读取 PVE 当前配置")
	}
	if err = s.exclusiveGPUMaintenance(ctx, db, d); err != nil {
		return d, "", err
	}
	blockers, _ := snapshotBaseBlockers(h, config)
	if len(blockers) > 0 {
		return d, "", errors.New(strings.Join(blockers, "；"))
	}
	if h.Digest != in.Digest {
		return d, "", errors.New("桌面配置已变化，请重新读取快照页面")
	}
	blocked, err := maintenanceBlocked(ctx, db, id, own, true)
	if err != nil {
		return d, "", err
	}
	if blocked {
		return d, "", errors.New("桌面有执行中任务或未确认维护结果")
	}
	var active bool
	if err = sqlx.GetContext(ctx, db, &active, "SELECT EXISTS(SELECT 1 FROM terminals WHERE session_desktop_id=$1 AND last_seen>now()-interval '90 seconds')", id); err != nil {
		return d, "", err
	}
	if active {
		return d, "", errors.New("桌面有活动连接，请先结束会话")
	}
	pve := s.currentPVE()
	permission := "VM.Snapshot"
	if in.Action == "rollback" {
		if in.ConfirmVMID != d.VMID {
			return d, "", errors.New("请填写当前 VMID，确认回退快照内磁盘数据和虚拟机配置")
		}
		ok, e := pve.HasPermission(ctx, fmt.Sprintf("/vms/%d", d.VMID), "VM.Snapshot.Rollback")
		if e != nil {
			return d, "", errors.New("无法核查快照还原权限")
		}
		if ok {
			permission = "VM.Snapshot.Rollback"
		}
	}
	ok, err := pve.HasPermission(ctx, fmt.Sprintf("/vms/%d", d.VMID), permission)
	if err != nil || !ok {
		return d, "", errors.New("缺少 " + permission + " 权限")
	}
	items, err := pve.Snapshots(ctx, d.Node, d.VMID)
	if err != nil {
		return d, "", errors.New("无法核对快照列表")
	}
	var selected *proxmox.Snapshot
	for i := range items {
		if items[i].Name == in.Name {
			selected = &items[i]
		}
	}
	if in.Action == "create" {
		if in.Revision != "" {
			return d, "", errors.New("创建快照不应携带已有快照版本")
		}
		if selected != nil {
			return d, "", errors.New("该快照名称已存在")
		}
		ok, err = pve.CanSnapshot(ctx, d.Node, d.VMID)
		if err != nil || !ok {
			return d, "", errors.New("PVE 当前磁盘或存储不支持快照")
		}
		return d, snapshotConfigHash(config), nil
	}
	if selected == nil || in.Revision == "" || in.Revision != snapshotRevision(*selected) {
		return d, "", errors.New("快照已删除或发生变化，请重新读取")
	}
	if selected.State != "" {
		return d, "", errors.New("快照状态未完成，请先在 PVE 核查")
	}
	if in.Action == "delete" {
		for _, item := range items {
			if item.Parent == in.Name {
				return d, "", errors.New("该快照仍有子快照，请先处理子快照")
			}
		}
	}
	target, err := pve.SnapshotConfig(ctx, d.Node, d.VMID, in.Name)
	if err != nil {
		return d, "", errors.New("无法读取快照配置")
	}
	if in.Action == "rollback" {
		if selected.Memory != 0 || fmt.Sprint(target["vmstate"]) != "<nil>" && fmt.Sprint(target["vmstate"]) != "" && fmt.Sprint(target["vmstate"]) != "0" {
			return d, "", errors.New("带内存状态的快照可能自动启动 VM，本期不支持还原")
		}
		currentDisks, e := proxmox.SnapshotDisks(config)
		if e != nil {
			return d, "", e
		}
		targetDisks, e := proxmox.SnapshotDisks(target)
		if e != nil {
			return d, "", e
		}
		a, _ := json.Marshal(currentDisks)
		b, _ := json.Marshal(targetDisks)
		if string(a) != string(b) {
			return d, "", errors.New("当前磁盘布局或容量与快照不同，请在 PVE 人工核查还原范围")
		}
		if proxmox.ConfigInt(target, "template", 0) != 0 {
			return d, "", errors.New("不能将桌面还原为模板")
		}
	}
	return d, snapshotConfigHash(target), nil
}
func (s *Server) checkSnapshot(c *gin.Context) {
	var in snapshotInput
	if !s.decodeConfiguration(c, &in) {
		return
	}
	_, _, err := s.snapshotPreflight(c.Request.Context(), s.DB, c.Param("id"), in, "")
	if err != nil {
		s.fail(c, 409, "snapshot_not_ready", err.Error())
		return
	}
	c.JSON(200, gin.H{"available": true, "message": "快照状态、桌面状态与权限预检通过"})
}
func (s *Server) queueSnapshot(c *gin.Context) {
	var in snapshotInput
	if !s.decodeConfiguration(c, &in) {
		return
	}
	key := strings.TrimSpace(c.GetHeader("Idempotency-Key"))
	if key == "" || len(key) > 100 {
		s.fail(c, 400, "idempotency_required", "请携带有效请求标识")
		return
	}
	tx, actor := s.accountTransaction(c, true)
	if tx == nil {
		return
	}
	defer tx.Rollback()
	if !s.lockAuthorization(c, tx) {
		return
	}
	ctx := c.Request.Context()
	raw, _ := json.Marshal(in)
	hash := store.Hash(string(raw))
	var old struct {
		TaskID string `db:"task_id"`
		Hash   string `db:"request_hash"`
	}
	err := tx.GetContext(ctx, &old, "SELECT j.task_id,j.request_hash FROM snapshot_jobs j JOIN tasks t ON t.id=j.task_id WHERE t.actor=$1 AND t.idempotency_key=$2", actor.Username, key)
	if err == nil {
		var task model.Task
		if old.Hash != hash {
			s.fail(c, 409, "idempotency_conflict", "请求标识已用于其他快照参数")
			return
		}
		if err = tx.GetContext(ctx, &task, "SELECT * FROM tasks WHERE id=$1 AND target_id=$2", old.TaskID, c.Param("id")); err != nil {
			s.fail(c, 409, "idempotency_conflict", "请求标识已用于其他桌面")
			return
		}
		c.JSON(202, task)
		return
	}
	if !errors.Is(err, sql.ErrNoRows) {
		s.dbError(c, err)
		return
	}
	var id string
	if err = tx.GetContext(ctx, &id, "SELECT id FROM desktops WHERE id=$1 FOR UPDATE", c.Param("id")); err != nil {
		s.dbError(c, err)
		return
	}
	d, targetHash, err := s.snapshotPreflight(ctx, tx, id, in, "")
	if err != nil {
		s.fail(c, 409, "snapshot_not_ready", err.Error())
		return
	}
	var task model.Task
	err = tx.GetContext(ctx, &task, "INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,message,idempotency_key) VALUES($1,'snapshot',$2,$3,$4,$5,'queued','已提交，执行前会再次预检',$6) RETURNING *", store.ID("task_"), "snapshot-"+in.Action, d.ID, d.Name, actor.Username, key)
	if conflict(err) {
		s.fail(c, 409, "task_conflict", "桌面已有执行中任务或请求标识已使用")
		return
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, "INSERT INTO snapshot_jobs(task_id,request_json,request_hash,target_hash) VALUES($1,$2,$3,$4)", task.ID, string(raw), hash, targetHash)
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, "INSERT INTO audit_logs(id,actor,action,target,result,request_id) VALUES($1,$2,$3,$4,'submitted',$5)", store.ID("audit_"), actor.Username, task.Action, d.ID, c.GetString("requestId"))
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(202, task)
}

type snapshotJob struct {
	Raw        string `db:"request_json"`
	Stage      string `db:"stage"`
	TargetHash string `db:"target_hash"`
}

func (s *Server) claimSnapshot(ctx context.Context, t model.Task, in snapshotInput) (bool, error) {
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		return false, err
	}
	defer tx.Rollback()
	var id string
	if err = tx.GetContext(ctx, &id, "SELECT id FROM desktops WHERE id=$1 FOR UPDATE", t.TargetID); err != nil {
		return false, err
	}
	result, err := tx.ExecContext(ctx, `UPDATE snapshot_jobs SET stage='submitting' WHERE task_id=$1 AND stage='prepared' AND EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND kind='snapshot' AND status IN ('queued','running'))`, t.ID)
	if err != nil {
		return false, err
	}
	n, _ := result.RowsAffected()
	if n != 1 {
		return false, nil
	}
	if in.Action == "rollback" {
		_, err = tx.ExecContext(ctx, "UPDATE desktops SET setup_required=true,service_state='unknown',updated_at=now() WHERE id=$1", t.TargetID)
		if err == nil {
			_, err = tx.ExecContext(ctx, "DELETE FROM connection_tickets WHERE desktop_id=$1", t.TargetID)
		}
	}
	if err == nil {
		err = tx.Commit()
	}
	return err == nil, err
}
func (s *Server) confirmSnapshot(ctx context.Context, db sqlx.ExtContext, t model.Task, in snapshotInput, targetHash string) error {
	if t.ProviderRef == "" {
		return errors.New("缺失原 PVE 任务标识，需人工核查")
	}
	pve := s.currentPVE()
	if pve == nil {
		return errors.New("计算平台未配置")
	}
	var d model.Desktop
	if err := sqlx.GetContext(ctx, db, &d, "SELECT * FROM desktops WHERE id=$1", t.TargetID); err != nil {
		return err
	}
	task, err := pve.Task(ctx, d.Node, t.ProviderRef)
	if err != nil {
		return err
	}
	if task.Status != "stopped" || task.ExitStatus != "OK" {
		return errors.New("原 PVE 任务未成功完成，请在 PVE 核查")
	}
	h, config, err := pve.Hardware(ctx, d.Node, d.VMID)
	if err != nil {
		return err
	}
	if h.PowerState != "stopped" || len(h.Blockers) > 0 {
		return errors.New("桌面状态仍未稳定，需核查")
	}
	items, err := pve.Snapshots(ctx, d.Node, d.VMID)
	if err != nil {
		return err
	}
	var found *proxmox.Snapshot
	for i := range items {
		if items[i].Name == in.Name {
			found = &items[i]
		}
	}
	switch in.Action {
	case "delete":
		if found != nil {
			return errors.New("快照尚未删除")
		}
	case "create":
		if found == nil || found.Memory != 0 || found.State != "" || found.Description != in.Description {
			return errors.New("新快照尚未确认")
		}
		snap, e := pve.SnapshotConfig(ctx, d.Node, d.VMID, in.Name)
		if e != nil {
			return e
		}
		if snapshotConfigHash(snap) != targetHash {
			return errors.New("新快照配置与创建前不一致")
		}
	case "rollback":
		if found == nil || found.Memory != 0 || found.State != "" || snapshotConfigHash(config) != targetHash {
			return errors.New("还原后的配置尚未确认")
		}
	default:
		return errors.New("未知快照操作")
	}
	return nil
}
func (s *Server) executeSnapshot(parent context.Context, t model.Task) {
	ctx, cancel := context.WithTimeout(parent, 30*time.Minute)
	defer cancel()
	var job snapshotJob
	err := s.DB.GetContext(ctx, &job, "SELECT request_json,stage,target_hash FROM snapshot_jobs WHERE task_id=$1", t.ID)
	var in snapshotInput
	if err == nil {
		err = json.Unmarshal([]byte(job.Raw), &in)
	}
	submitted := job.Stage != "prepared"
	if e := s.RefreshCompute(ctx); err == nil && e != nil {
		err = e
	}
	if err == nil && !submitted {
		var d model.Desktop
		var targetHash string
		d, targetHash, err = s.snapshotPreflight(ctx, s.DB, t.TargetID, in, t.ID)
		if err == nil && targetHash != job.TargetHash {
			err = errors.New("快照配置已变化")
		}
		if err == nil {
			var claimed bool
			claimed, err = s.claimSnapshot(ctx, t, in)
			if err == nil && !claimed {
				return
			}
			if err == nil {
				submitted = true
				t.ProviderRef, err = s.currentPVE().SnapshotAction(ctx, d.Node, d.VMID, in.Action, in.Name, in.Description)
				var rejected *proxmox.HTTPError
				if errors.As(err, &rejected) && rejected.Status >= 400 && rejected.Status < 500 {
					submitted = false
				}
				if err == nil {
					_, err = s.DB.ExecContext(ctx, "UPDATE tasks SET provider_ref=$1,updated_at=now() WHERE id=$2", t.ProviderRef, t.ID)
				}
			}
		}
	}
	if err == nil && t.ProviderRef == "" {
		err = errors.New("原 PVE 任务标识缺失")
	}
	if err == nil {
		var d model.Desktop
		err = s.DB.GetContext(ctx, &d, "SELECT * FROM desktops WHERE id=$1", t.TargetID)
		if err == nil {
			err = s.currentPVE().Wait(ctx, d.Node, t.ProviderRef, time.Second)
		}
	}
	if err == nil {
		err = s.confirmSnapshot(ctx, s.DB, t, in, job.TargetHash)
	}
	status, stage, message := "success", "completed", "PVE 快照任务及实际状态已确认"
	if in.Action == "rollback" {
		message = "快照还原已确认，桌面保持关机；请重新确认 IP、系统身份、Sunshine 服务与终端配对后开放连接"
	}
	if err != nil {
		status, stage, message = "failed", "failed", "快照操作未提交，请重新检查关机状态、快照版本、权限和存储能力"
		if submitted {
			status, stage, message = "timeout", "uncertain", "快照操作结果未确认，已阻止后续桌面修改和新连接；仅核对原任务，不会自动重发"
		}
	}
	s.finishSnapshot(t, status, stage, message)
}
func (s *Server) finishSnapshot(t model.Task, status, stage, message string) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		return
	}
	defer tx.Rollback()
	var id, current, jobStage string
	if err = tx.GetContext(ctx, &id, "SELECT id FROM desktops WHERE id=$1 FOR UPDATE", t.TargetID); err != nil {
		return
	}
	if err = tx.GetContext(ctx, &current, "SELECT status FROM tasks WHERE id=$1 FOR UPDATE", t.ID); err != nil || (current != "queued" && current != "running") {
		return
	}
	if err = tx.GetContext(ctx, &jobStage, "SELECT stage FROM snapshot_jobs WHERE task_id=$1 FOR UPDATE", t.ID); err != nil {
		return
	}
	if stage == "expired" {
		status, stage, message = "failed", "failed", "快照任务提交前过期，未发送修改"
		if jobStage != "prepared" {
			status, stage, message = "timeout", "uncertain", "快照任务结果未确认，请核对原任务，不会自动重发"
		}
	}
	if _, err = tx.ExecContext(ctx, "UPDATE snapshot_jobs SET stage=$1 WHERE task_id=$2", stage, t.ID); err == nil {
		err = s.completeTx(ctx, tx, t, status, message)
	}
	if err == nil {
		_ = tx.Commit()
	}
}
func (s *Server) reconcileSnapshot(c *gin.Context) {
	tx, actor := s.accountTransaction(c, true)
	if tx == nil {
		return
	}
	defer tx.Rollback()
	ctx := c.Request.Context()
	var t model.Task
	if err := tx.GetContext(ctx, &t, "SELECT * FROM tasks WHERE id=$1 AND kind='snapshot'", c.Param("id")); err != nil {
		s.dbError(c, err)
		return
	}
	var id string
	if err := tx.GetContext(ctx, &id, "SELECT id FROM desktops WHERE id=$1 FOR UPDATE", t.TargetID); err != nil {
		s.dbError(c, err)
		return
	}
	if err := tx.GetContext(ctx, &t, "SELECT * FROM tasks WHERE id=$1 FOR UPDATE", t.ID); err != nil {
		s.dbError(c, err)
		return
	}
	var job snapshotJob
	if err := tx.GetContext(ctx, &job, "SELECT request_json,stage,target_hash FROM snapshot_jobs WHERE task_id=$1 FOR UPDATE", t.ID); err != nil {
		s.dbError(c, err)
		return
	}
	if t.Status != "timeout" || job.Stage != "uncertain" {
		s.fail(c, 409, "reconcile_unavailable", "仅可核对结果未确认的快照任务")
		return
	}
	var in snapshotInput
	if json.Unmarshal([]byte(job.Raw), &in) != nil {
		s.fail(c, 500, "invalid_job", "任务记录无法读取")
		return
	}
	if err := s.confirmSnapshot(ctx, tx, t, in, job.TargetHash); err != nil {
		s.fail(c, 409, "snapshot_unconfirmed", "原任务或实际状态尚未确认，请在 PVE 核查；未重新发送修改")
		return
	}
	_, err := tx.ExecContext(ctx, "UPDATE snapshot_jobs SET stage='completed' WHERE task_id=$1", t.ID)
	if err == nil {
		_, err = tx.ExecContext(ctx, "UPDATE tasks SET status='success',message='只读核对原 PVE 任务和实际状态通过，未重新发送修改',updated_at=now() WHERE id=$1", t.ID)
	}
	if err == nil {
		err = accountAudit(ctx, tx, actor.Username, "reconcile_snapshot", t.ID, c.GetString("requestId"))
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(200, gin.H{"accepted": true})
}