feat: add full snapshot downloads and safe manual restore

2026-10-17 01:36:21 +08:00
parent c6615faa49
commit f77f3d77c0
12 changed files with 1318 additions and 7 deletions
@@ -0,0 +1,273 @@
+package api
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"context"
+	"errors"
+	"io"
+	"mime"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strings"
+	"syscall"
+	"testing"
+)
+
+type snapshotTransferWriter struct {
+	*httptest.ResponseRecorder
+	check func()
+}
+
+func (w *snapshotTransferWriter) WriteHeader(code int) {
+	w.check()
+	w.ResponseRecorder.WriteHeader(code)
+}
+
+func TestSnapshotDownloadTransferUnlockedAndOversize(t *testing.T) {
+	a, i, home := configTestAPI(t, false)
+	tmp := t.TempDir()
+	t.Setenv("TMPDIR", tmp)
+	id := "failed-home-" + strings.Repeat("c", 32)
+	dir := filepath.Join(filepath.Dir(home), id)
+	if err := os.Mkdir(dir, 0700); err != nil {
+		t.Fatal(err)
+	}
+	configWrite(t, filepath.Join(dir, "data"), "saved")
+	checked := false
+	w := &snapshotTransferWriter{httptest.NewRecorder(), func() {
+		checked = true
+		if a.instanceLocks.held() {
+			t.Fatal("instance locked during transfer")
+		}
+		files, err := os.ReadDir(tmp)
+		if err != nil || len(files) != 1 {
+			t.Fatal(files, err)
+		}
+		st, err := files[0].Info()
+		if err != nil || st.Mode().Perm() != 0600 {
+			t.Fatal(st, err)
+		}
+		if err := os.RemoveAll(dir); err != nil {
+			t.Fatal(err)
+		}
+	}}
+	a.snapshotDownload(w, httptest.NewRequest(http.MethodGet, "/?download="+id, nil), i)
+	if !checked || w.Code != 200 {
+		t.Fatal(checked, w.Code)
+	}
+	gz, err := gzip.NewReader(bytes.NewReader(w.Body.Bytes()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer gz.Close()
+	if _, err := io.Copy(io.Discard, gz); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Mkdir(dir, 0700); err != nil {
+		t.Fatal(err)
+	}
+	f, err := os.Create(filepath.Join(dir, "huge"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = f.Truncate(upgradeSnapshotLimit + 1)
+	f.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+	rec := httptest.NewRecorder()
+	a.snapshotDownload(rec, httptest.NewRequest("GET", "/?download="+id, nil), i)
+	if rec.Code != 413 || !strings.Contains(rec.Body.String(), "limit") {
+		t.Fatal(rec.Code, rec.Body.String())
+	}
+	files, err := os.ReadDir(tmp)
+	if err != nil || len(files) != 0 {
+		t.Fatal(files, err)
+	}
+}
+
+func TestSnapshotDownloadContentsAndCleanup(t *testing.T) {
+	a, i, home := configTestAPI(t, true)
+	tmp := t.TempDir()
+	t.Setenv("TMPDIR", tmp)
+	id := "pre-restore-" + strings.Repeat("a", 32)
+	dir := filepath.Join(filepath.Dir(home), id)
+	if err := os.MkdirAll(filepath.Join(dir, "node_modules/.pnpm/pkg"), 0750); err != nil {
+		t.Fatal(err)
+	}
+	configWrite(t, filepath.Join(dir, "node_modules/.pnpm/pkg/bin"), "#!/bin/sh")
+	if err := os.Chmod(filepath.Join(dir, "node_modules/.pnpm/pkg/bin"), 0751); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Symlink(".pnpm/pkg", filepath.Join(dir, "node_modules/pkg")); err != nil {
+		t.Fatal(err)
+	}
+	configWrite(t, filepath.Join(dir, ".credentials.yaml"), "secret")
+	if _, err := a.s.DB.Exec("UPDATE instances SET status='error' WHERE id=?", i.ID); err != nil {
+		t.Fatal(err)
+	}
+	w := httptest.NewRecorder()
+	a.config(w, httptest.NewRequest("GET", "/?download="+id, nil), i, "snapshots")
+	if w.Code != 200 {
+		t.Fatal(w.Code, w.Body.String())
+	}
+	_, params, err := mime.ParseMediaType(w.Header().Get("Content-Disposition"))
+	if err != nil || params["filename"] != id+".tar.gz" {
+		t.Fatal(params, err)
+	}
+	gz, err := gzip.NewReader(bytes.NewReader(w.Body.Bytes()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer gz.Close()
+	tr := tar.NewReader(gz)
+	seen := map[string]bool{}
+	for {
+		h, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+		seen[h.Name] = true
+		b, err := io.ReadAll(tr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		switch h.Name {
+		case "node_modules/.pnpm/pkg/bin":
+			if string(b) != "#!/bin/sh" || h.Mode != 0751 {
+				t.Fatal(h, string(b))
+			}
+		case "node_modules/pkg":
+			if h.Typeflag != tar.TypeSymlink || h.Linkname != ".pnpm/pkg" {
+				t.Fatal(h)
+			}
+		case ".credentials.yaml":
+			if string(b) != "secret" {
+				t.Fatal(string(b))
+			}
+		}
+	}
+	if len(seen) != 7 || !seen["node_modules/pkg"] || seen["node_modules/pkg/bin"] {
+		t.Fatal(seen)
+	}
+	files, err := os.ReadDir(tmp)
+	if err != nil || len(files) != 0 {
+		t.Fatal(files, err)
+	}
+	r := httptest.NewRequest("GET", "/?download="+id, nil)
+	r.Header.Set("Range", "bytes=0-9")
+	w = httptest.NewRecorder()
+	a.snapshotDownload(w, r, i)
+	if w.Code != 206 || w.Body.Len() != 10 {
+		t.Fatal(w.Code, w.Body.Len())
+	}
+}
+
+func TestSnapshotDownloadRejectsUnsafeAndForeign(t *testing.T) {
+	a, i, home := configTestAPI(t, false)
+	id := "pre-restore-" + strings.Repeat("b", 32)
+	dir := filepath.Join(filepath.Dir(home), id)
+	if err := os.Mkdir(dir, 0700); err != nil {
+		t.Fatal(err)
+	}
+	for _, query := range []string{"", "../home", "/home", id + "/../home", "pre-restore-ABC"} {
+		w := httptest.NewRecorder()
+		a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+url.QueryEscape(query), nil), i)
+		if w.Code != 400 {
+			t.Fatal(query, w.Code)
+		}
+	}
+	foreign := i
+	foreign.UID = "someone-else"
+	w := httptest.NewRecorder()
+	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), foreign)
+	if w.Code != 404 {
+		t.Fatal(w.Code)
+	}
+	for _, target := range []string{"../home", "/etc/passwd", "../../elsewhere"} {
+		if err := os.Symlink(target, filepath.Join(dir, "link")); err != nil {
+			t.Fatal(err)
+		}
+		w = httptest.NewRecorder()
+		a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
+		if w.Code != 400 || w.Header().Get("Content-Disposition") != "" {
+			t.Fatal(target, w.Code)
+		}
+		if err := os.Remove(filepath.Join(dir, "link")); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := syscall.Mkfifo(filepath.Join(dir, "fifo"), 0600); err != nil {
+		t.Fatal(err)
+	}
+	w = httptest.NewRecorder()
+	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
+	if w.Code != 400 {
+		t.Fatal(w.Code)
+	}
+	if err := os.Remove(filepath.Join(dir, "fifo")); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Remove(dir); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Symlink(home, dir); err != nil {
+		t.Fatal(err)
+	}
+	w = httptest.NewRecorder()
+	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
+	if w.Code != 400 {
+		t.Fatal(w.Code)
+	}
+}
+
+func TestSnapshotArchiveLimitsAndCancellation(t *testing.T) {
+	dir := t.TempDir()
+	configWrite(t, filepath.Join(dir, "file"), strings.Repeat("x", 2048))
+	root, err := os.OpenRoot(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer root.Close()
+	if err := snapshotArchive(context.Background(), root, io.Discard, 1024); !errors.Is(err, errSnapshotDownloadLimit) {
+		t.Fatal(err)
+	}
+	if err := snapshotArchive(context.Background(), root, io.Discard, 2048); err != nil {
+		t.Fatal("exact content limit should allow archive overhead", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := snapshotArchive(ctx, root, io.Discard, upgradeSnapshotLimit); !errors.Is(err, context.Canceled) {
+		t.Fatal(err)
+	}
+	a, i, _ := configTestAPI(t, false)
+	tmp := t.TempDir()
+	t.Setenv("TMPDIR", tmp)
+	r := httptest.NewRequest("GET", "/?download=pre-restore-"+strings.Repeat("a", 32), nil).WithContext(ctx)
+	w := httptest.NewRecorder()
+	a.snapshotDownload(w, r, i)
+	if w.Code != 408 {
+		t.Fatal(w.Code)
+	}
+	files, err := os.ReadDir(tmp)
+	if err != nil || len(files) != 0 {
+		t.Fatal(files, err)
+	}
+	if len(zipSlots) != 0 || a.instanceLocks.held() {
+		t.Fatal("resource leak")
+	}
+	t.Setenv("TMPDIR", filepath.Join(tmp, "missing"))
+	w = httptest.NewRecorder()
+	a.snapshotDownload(w, httptest.NewRequest("GET", r.URL.String(), nil), i)
+	if w.Code != 500 {
+		t.Fatal(w.Code)
+	}
+}