dsh-web/internal/api/snapshot_restore.go

package api

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"git.misaka.ren/m1saka/dsh-web/internal/store"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const snapshotUnknownReason = "snapshot source image/version unknown; explicit acknowledgement required"

func snapshotRuntimeLayout(config, host map[string]json.RawMessage, mounts []struct{ Type, Source, Destination string }, home string) bool {
	workspace := 0
	for _, mount := range mounts {
		if mount.Type == "tmpfs" && mount.Destination == "/tmp" {
			continue
		}
		if mount.Type != "bind" || mount.Source != home || mount.Destination != "/workspace" {
			return false
		}
		workspace++
	}
	if workspace != 1 {
		return false
	}
	var env []string
	if json.Unmarshal(config["Env"], &env) != nil {
		return false
	}
	values := map[string]string{}
	for _, entry := range env {
		key, value, _ := strings.Cut(entry, "=")
		values[key] = value
	}
	for key, want := range map[string]string{"DSH_HOME": "/workspace", "HOME": "/workspace", "XDG_DATA_HOME": "/workspace/.local/share", "XDG_CONFIG_HOME": "/workspace/.config"} {
		if values[key] != want {
			return false
		}
	}
	for index, raw := range []json.RawMessage{host["Tmpfs"], config["Volumes"]} {
		if len(raw) == 0 {
			continue
		}
		var mounts map[string]json.RawMessage
		if json.Unmarshal(raw, &mounts) != nil {
			return false
		}
		if index == 0 {
			delete(mounts, "/tmp")
		}
		if len(mounts) != 0 {
			return false
		}
	}
	return true
}

func (a *API) recordSnapshot(i store.Instance, id, version, image string) error {
	_, err := a.s.DB.Exec("CREATE TABLE IF NOT EXISTS snapshot_metadata (instance_id TEXT NOT NULL, id TEXT NOT NULL, source_version TEXT NOT NULL, source_image TEXT NOT NULL, PRIMARY KEY(instance_id,id))")
	if err != nil {
		return err
	}
	_, err = a.s.DB.Exec("INSERT INTO snapshot_metadata(instance_id,id,source_version,source_image) VALUES(?,?,?,?)", i.ID, id, version, image)
	return err
}

func (a *API) snapshotCompatibility(i store.Instance, id string) (version, image, reason string) {
	var exists int
	if err := a.s.DB.QueryRow("SELECT count(*) FROM sqlite_master WHERE type='table' AND name='snapshot_metadata'").Scan(&exists); err != nil {
		return "", "", "snapshot metadata unavailable"
	}
	if exists != 0 {
		err := a.s.DB.QueryRow("SELECT source_version,source_image FROM snapshot_metadata WHERE instance_id=? AND id=?", i.ID, id).Scan(&version, &image)
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			return "", "", "snapshot metadata unavailable"
		}
	}
	if (version != "" && version != i.Version) || (image != "" && image != i.ImageID) {
		return version, image, "snapshot source image/version is incompatible"
	}
	if version == "" || image == "" {
		return version, image, snapshotUnknownReason
	}
	return version, image, ""
}

func (a *API) snapshotRestore(w http.ResponseWriter, r *http.Request, i store.Instance) {
	controller := http.NewResponseController(w)
	w, flush := deferJSONResponse(w)
	defer flush()
	select {
	case configBodySlots <- struct{}{}:
		defer func() { r.Body = http.NoBody; <-configBodySlots }()
	default:
		fail(w, 503, "configuration uploads busy; retry later")
		return
	}
	if !configReadBody(w, r, controller, 4096) {
		return
	}
	var in struct {
		ID           string
		Password     string
		Confirm      bool
		AllowUnknown bool `json:"allow_unknown_version"`
	}
	dec := json.NewDecoder(r.Body)
	if dec.Decode(&in) != nil || dec.Decode(new(any)) != io.EOF || !snapshotIDPattern.MatchString(in.ID) || !in.Confirm {
		fail(w, 400, "valid snapshot ID and explicit confirmation required")
		return
	}
	lock := a.instanceLocks.Acquire(i.ID)
	defer lock.Release()
	a.readersMu.Lock()
	readers := a.workspaceReaders[i.ID]
	a.readersMu.Unlock()
	if readers > 0 {
		fail(w, 409, "instance files busy; retry later")
		return
	}
	user := a.userLocks.Acquire(i.UID)
	defer user.Release()
	current, err := a.s.GetInstance(i.ID)
	if err != nil || current.UID != i.UID {
		fail(w, 404, "not found")
		return
	}
	i = current
	if !a.reauthenticate(w, i.UID, in.Password) {
		return
	}
	if reason := a.backupProtection(i); reason != "" {
		fail(w, 409, reason)
		return
	}
	admissionMu.RLock()
	defer admissionMu.RUnlock()
	buildQueue.Lock()
	err = a.ensureBuilds()
	buildQueue.Unlock()
	if err != nil {
		fail(w, 500, "database error")
		return
	}
	var blocked int
	if err = a.s.DB.QueryRowContext(r.Context(), "SELECT count(*) FROM image_admissions WHERE image_id=? AND state='blocked'", i.ImageID).Scan(&blocked); err != nil {
		fail(w, 500, "database error")
		return
	}
	if blocked > 0 {
		fail(w, 409, "image is blocked; restore denied")
		return
	}
	version, image, reason := a.snapshotCompatibility(i, in.ID)
	if reason != "" && !(reason == snapshotUnknownReason && in.AllowUnknown) {
		fail(w, 409, reason)
		return
	}
	root, err := managementRoot(a.root, "workspaces", i.UID, "instances", i.ID)
	if errors.Is(err, os.ErrNotExist) {
		fail(w, 404, "home or snapshot not found")
		return
	}
	if err != nil {
		managementPathError(w, err)
		return
	}
	defer root.Close()
	for _, name := range []string{"home", in.ID} {
		st, e := root.Lstat(name)
		if errors.Is(e, os.ErrNotExist) {
			fail(w, 404, "home or snapshot not found")
			return
		}
		if e != nil {
			managementPathError(w, e)
			return
		}
		if !st.IsDir() {
			fail(w, 400, "unsafe home or snapshot")
			return
		}
	}
	f, err := root.Open(".")
	if err != nil {
		fail(w, 500, "snapshot storage unavailable")
		return
	}
	entries, err := f.ReadDir(-1)
	f.Close()
	count := 0
	for _, entry := range entries {
		if strings.HasPrefix(entry.Name(), "pre-restore-") || strings.HasPrefix(entry.Name(), "failed-restore-") {
			count++
		}
	}
	if err != nil {
		fail(w, 500, "snapshot listing failed")
		return
	}
	if count >= 3 {
		fail(w, 409, "restore snapshot limit reached; cleanup required")
		return
	}
	size, err := quotaTreeBytes(root, in.ID)
	var pathError *os.PathError
	if errors.As(err, &pathError) {
		managementPathError(w, err)
		return
	}
	if err != nil || a.quotaGrowth(i.UID, size) != nil {
		fail(w, 413, "quota exceeded or snapshot size unavailable")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 3*time.Minute)
	defer cancel()
	var container struct {
		Image              string
		Config, HostConfig map[string]json.RawMessage
		Mounts             []struct{ Type, Source, Destination string }
		State              struct{ Running, Paused, Restarting bool }
	}
	err = a.docker.Do(ctx, "GET", "/containers/"+url.PathEscape(i.ContainerID)+"/json", nil, &container)
	if err != nil {
		fail(w, 502, "container inspection failed")
		return
	}
	if container.State.Paused || container.State.Restarting {
		fail(w, 409, "unsupported container state")
		return
	}
	var runtimeImage struct {
		ID     string `json:"Id"`
		Config struct{ Labels map[string]string }
	}
	hostRoot := os.Getenv("HOST_DATA_ROOT")
	if hostRoot == "" {
		hostRoot = a.root
	}
	homeSource, homeErr := toHostPath(a.root, hostRoot, filepath.Join(root.Name(), "home"))
	if homeErr != nil || !snapshotRuntimeLayout(container.Config, container.HostConfig, container.Mounts, homeSource) {
		fail(w, 409, "incompatible effective runtime workspace layout")
		return
	}
	if i.ImageID == "" || container.Image != i.ImageID || container.Config == nil || container.HostConfig == nil || homeErr != nil || a.docker.Do(ctx, "GET", "/images/"+url.PathEscape(i.ImageID)+"/json", nil, &runtimeImage) != nil || runtimeImage.ID != i.ImageID || runtimeImage.Config.Labels["dsh-mgr.version"] != i.Version || upgradeWorkspace(container.Config, container.HostConfig, runtimeImage.Config.Labels["dsh-mgr.recipe"], homeSource) != nil {
		fail(w, 409, "incompatible runtime image or workspace layout")
		return
	}
	wasRunning := container.State.Running
	op := store.ID()
	staging := ".snapshot-restore-" + op
	previous := "pre-restore-" + op
	failed := "failed-restore-" + op
	if err = upgradeCopy(ctx, root, in.ID, staging, size); err != nil {
		if root.RemoveAll(staging) != nil {
			a.s.SetInstance(i.ID, "error")
			fail(w, 500, "staging cleanup requires administrator recovery")
			return
		}
		if errors.As(err, &pathError) {
			managementPathError(w, err)
			return
		}
		fail(w, 400, "unsafe or oversized snapshot")
		return
	}
	defer func() {
		if root.RemoveAll(staging) != nil {
			a.s.SetInstance(i.ID, "error")
		}
	}()
	if err = a.recordSnapshot(i, previous, i.Version, i.ImageID); err != nil {
		fail(w, 500, "snapshot metadata unavailable")
		return
	}
	if err = a.recordSnapshot(i, failed, version, image); err != nil {
		fail(w, 500, "snapshot metadata unavailable")
		return
	}
	if _, err = a.s.DB.Exec("UPDATE instances SET status='error' WHERE id=?", i.ID); err != nil {
		fail(w, 500, "cannot record restore state")
		return
	}
	if _, err = a.s.DB.Exec("INSERT INTO audit(uid,action,target,detail,at) VALUES(?,?,?,?,?)", i.UID, "snapshot.restore.started", i.ID, in.ID+" "+previous+" "+failed, time.Now().Unix()); err != nil {
		fail(w, 500, "restore audit unavailable; recovery required")
		return
	}
	if err = a.configStop(ctx, i.ContainerID); err != nil {
		fail(w, 502, "runtime stop failed; recovery required")
		return
	}
	moved, installed := false, false
	var owner store.User
	owner.ID = i.UID
	err = a.s.DB.QueryRow("SELECT quota FROM users WHERE id=?", i.UID).Scan(&owner.Quota)
	if err == nil {
		var used int64
		used, err = a.userUsage(owner)
		if err == nil && used > owner.Quota {
			err = errors.New("quota exceeded after runtime stop")
		}
	}
	if err == nil {
		err = root.Rename("home", previous)
	}
	if err == nil {
		moved = true
		err = upgradeSync(root)
	}
	if err == nil {
		err = root.Rename(staging, "home")
		installed = err == nil
	}
	if err == nil {
		err = upgradeSync(root)
	}
	startReady := func(c context.Context) error {
		runtimeAuthCache.Delete(runtimeAuthKey{a, i.ID})
		if e := a.docker.Do(c, "POST", "/containers/"+url.PathEscape(i.ContainerID)+"/start", nil, nil); e != nil {
			return e
		}
		return a.upgradeReady(c, i.ContainerID, i.ID, a.runtimeHost(r))
	}
	if err == nil {
		err = startReady(ctx)
	}
	if err == nil && !wasRunning {
		err = a.configStop(ctx, i.ContainerID)
	}
	rejected := err != nil
	if rejected {
		rollback, done := context.WithTimeout(context.Background(), 90*time.Second)
		defer done()
		if a.configStop(rollback, i.ContainerID) != nil {
			fail(w, 500, "restore rollback requires runtime stop; recovery required")
			return
		}
		if installed {
			if root.Rename("home", failed) != nil || upgradeSync(root) != nil {
				fail(w, 500, "restore rollback requires administrator recovery")
				return
			}
		}
		if moved {
			if root.Rename(previous, "home") != nil || upgradeSync(root) != nil {
				fail(w, 500, "restore rollback requires administrator recovery")
				return
			}
		}
		if wasRunning {
			if startReady(rollback) != nil {
				fail(w, 500, "restore rollback readiness failed; recovery required")
				return
			}
		}
	}
	status := "exited"
	if wasRunning {
		status = "running"
	}
	if _, err = a.s.DB.Exec("UPDATE instances SET status=? WHERE id=?", status, i.ID); err != nil {
		fail(w, 500, "restore state requires administrator recovery")
		return
	}
	runtimeAuthCache.Delete(runtimeAuthKey{a, i.ID})
	if rejected {
		a.s.AddAudit(i.UID, "snapshot.restore.rolled_back", i.ID, in.ID)
		fail(w, 400, "restore rejected; previous home restored")
		return
	}
	a.s.AddAudit(i.UID, "snapshot.restore.completed", i.ID, in.ID)
	writeJSON(w, map[string]any{"ok": true, "snapshot_id": previous})
}