nekonest-cloud/relay/internal/tenantbackup/backup.go

package tenantbackup

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/klarkxy/nekonest-cloud/relay/internal/tenantfs"
	_ "modernc.org/sqlite"
)

const manifestVersion = 1

type File struct {
	Path   string `json:"path"`
	Size   int64  `json:"size"`
	SHA256 string `json:"sha256"`
}

type Manifest struct {
	Version             int    `json:"version"`
	TenantID            string `json:"tenant_id"`
	PlacementGeneration int64  `json:"placement_generation"`
	CreatedAt           string `json:"created_at"`
	Files               []File `json:"files"`
}

type restoreReceipt struct {
	Version             int    `json:"version"`
	TenantID            string `json:"tenant_id"`
	PlacementGeneration int64  `json:"placement_generation"`
	ManifestSHA256      string `json:"manifest_sha256"`
}

type Result struct {
	Path           string
	BackupRef      string
	Manifest       Manifest
	ManifestSHA256 string
}

func validBackupRef(value string) bool {
	parts := strings.Split(filepath.ToSlash(value), "/")
	if len(parts) != 2 || len(parts[0]) != 32 || !strings.HasPrefix(parts[1], "g") || strings.Contains(parts[1], ".tmp") {
		return false
	}
	for _, char := range parts[0] {
		if !strings.ContainsRune("0123456789abcdef", char) {
			return false
		}
	}
	if len(parts[1]) != len("g00000000000000000000-20060102T150405Z-0000000000000000") {
		return false
	}
	for index, char := range parts[1] {
		switch index {
		case 0:
			if char != 'g' {
				return false
			}
		case 21, 38:
			if char != '-' {
				return false
			}
		case 30:
			if char != 'T' {
				return false
			}
		case 37:
			if char != 'Z' {
				return false
			}
		default:
			if index >= 39 {
				if !strings.ContainsRune("0123456789abcdef", char) {
					return false
				}
			} else if char < '0' || char > '9' {
				return false
			}
		}
	}
	return true
}

// ResolveReference turns an opaque control-plane backup reference into a
// local immutable backup path without accepting an absolute path or symlink.
func ResolveReference(backupRoot, reference string) (string, error) {
	if !validBackupRef(reference) {
		return "", fmt.Errorf("invalid backup reference")
	}
	root, err := filepath.Abs(strings.TrimSpace(backupRoot))
	if err != nil || strings.TrimSpace(backupRoot) == "" {
		return "", fmt.Errorf("invalid backup root")
	}
	if info, err := os.Lstat(root); err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return "", fmt.Errorf("backup root is not a real directory")
	}
	parts := strings.Split(filepath.ToSlash(reference), "/")
	parent := filepath.Join(root, parts[0])
	if info, err := os.Lstat(parent); err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return "", fmt.Errorf("backup tenant directory is not real")
	}
	resolved := filepath.Join(parent, parts[1])
	if filepath.Dir(resolved) != parent {
		return "", fmt.Errorf("backup reference escaped root")
	}
	if info, err := os.Lstat(resolved); err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return "", fmt.Errorf("backup reference is unavailable")
	}
	return resolved, nil
}

func randomHex(size int) (string, error) {
	value := make([]byte, size)
	if _, err := rand.Read(value); err != nil {
		return "", err
	}
	return hex.EncodeToString(value), nil
}

func ensureRealDirectory(path string) error {
	if err := os.MkdirAll(path, 0o700); err != nil {
		return err
	}
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return fmt.Errorf("%s is not a real directory", path)
	}
	return os.Chmod(path, 0o700)
}

func regularFile(path string) (os.FileInfo, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return nil, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return nil, fmt.Errorf("%s is not a regular file", path)
	}
	return info, nil
}

func validAttachmentName(name string) bool {
	stem := ""
	switch {
	case strings.HasSuffix(name, ".bin"):
		stem = strings.TrimSuffix(name, ".bin")
	case strings.HasSuffix(name, ".json"):
		stem = strings.TrimSuffix(name, ".json")
	default:
		return false
	}
	if len(stem) != 32 {
		return false
	}
	for _, char := range stem {
		if !strings.ContainsRune("0123456789abcdef", char) {
			return false
		}
	}
	return true
}

func checkpointAndVerify(path string) error {
	if _, err := regularFile(path); err != nil {
		return fmt.Errorf("inspect sqlite database: %w", err)
	}
	database, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)")
	if err != nil {
		return err
	}
	defer database.Close()
	var busy, logFrames, checkpointed int
	if err := database.QueryRow(`PRAGMA wal_checkpoint(TRUNCATE)`).Scan(&busy, &logFrames, &checkpointed); err != nil {
		return fmt.Errorf("checkpoint sqlite: %w", err)
	}
	if busy != 0 {
		return fmt.Errorf("sqlite checkpoint remained busy")
	}
	var result string
	if err := database.QueryRow(`PRAGMA integrity_check`).Scan(&result); err != nil {
		return fmt.Errorf("check sqlite integrity: %w", err)
	}
	if result != "ok" {
		return fmt.Errorf("sqlite integrity check failed: %s", result)
	}
	return nil
}

func verifySQLiteReadOnly(path string) error {
	if _, err := regularFile(path); err != nil {
		return fmt.Errorf("inspect sqlite database: %w", err)
	}
	database, err := sql.Open("sqlite", "file:"+filepath.ToSlash(path)+"?mode=ro&immutable=1&_pragma=query_only(1)")
	if err != nil {
		return err
	}
	defer database.Close()
	var result string
	if err := database.QueryRow(`PRAGMA integrity_check`).Scan(&result); err != nil {
		return fmt.Errorf("check restored sqlite integrity: %w", err)
	}
	if result != "ok" {
		return fmt.Errorf("restored sqlite integrity check failed: %s", result)
	}
	return nil
}

func hashFile(ctx context.Context, path string) (int64, string, error) {
	if _, err := regularFile(path); err != nil {
		return 0, "", err
	}
	file, err := os.Open(path)
	if err != nil {
		return 0, "", err
	}
	defer file.Close()
	hash := sha256.New()
	buffer := make([]byte, 128<<10)
	var total int64
	for {
		if err := ctx.Err(); err != nil {
			return 0, "", err
		}
		count, readErr := file.Read(buffer)
		if count > 0 {
			total += int64(count)
			_, _ = hash.Write(buffer[:count])
		}
		if errors.Is(readErr, io.EOF) {
			break
		}
		if readErr != nil {
			return 0, "", readErr
		}
	}
	return total, hex.EncodeToString(hash.Sum(nil)), nil
}

func copyFile(ctx context.Context, source, destination string) (File, error) {
	if _, err := regularFile(source); err != nil {
		return File{}, err
	}
	input, err := os.Open(source)
	if err != nil {
		return File{}, err
	}
	defer input.Close()
	if err := ensureRealDirectory(filepath.Dir(destination)); err != nil {
		return File{}, err
	}
	output, err := os.OpenFile(destination, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return File{}, err
	}
	remove := true
	defer func() {
		_ = output.Close()
		if remove {
			_ = os.Remove(destination)
		}
	}()
	hash := sha256.New()
	written, err := io.Copy(io.MultiWriter(output, hash), &contextReader{ctx: ctx, reader: input})
	if err != nil {
		return File{}, err
	}
	if err := output.Sync(); err != nil {
		return File{}, err
	}
	if err := output.Close(); err != nil {
		return File{}, err
	}
	remove = false
	return File{Size: written, SHA256: hex.EncodeToString(hash.Sum(nil))}, nil
}

type contextReader struct {
	ctx    context.Context
	reader io.Reader
}

func (reader *contextReader) Read(buffer []byte) (int, error) {
	if err := reader.ctx.Err(); err != nil {
		return 0, err
	}
	return reader.reader.Read(buffer)
}

func sourceFiles(paths tenantfs.Paths) ([]struct{ absolute, relative string }, error) {
	if _, err := regularFile(paths.Database); err != nil {
		return nil, err
	}
	attachmentInfo, err := os.Lstat(paths.Attachments)
	if err != nil {
		return nil, err
	}
	if attachmentInfo.Mode()&os.ModeSymlink != 0 || !attachmentInfo.IsDir() {
		return nil, fmt.Errorf("attachment root is not a real directory")
	}
	result := []struct{ absolute, relative string }{{paths.Database, "relay.db"}}
	entries, err := os.ReadDir(paths.Attachments)
	if err != nil {
		return nil, err
	}
	for _, entry := range entries {
		if entry.IsDir() || entry.Type()&os.ModeSymlink != 0 || !validAttachmentName(entry.Name()) {
			return nil, fmt.Errorf("invalid attachment artifact %q", entry.Name())
		}
		result = append(result, struct{ absolute, relative string }{
			filepath.Join(paths.Attachments, entry.Name()), filepath.ToSlash(filepath.Join("attachments", entry.Name())),
		})
	}
	sort.Slice(result, func(left, right int) bool { return result[left].relative < result[right].relative })
	return result, nil
}

func Create(ctx context.Context, dataRoot, backupRoot, tenantID string, generation int64, now time.Time) (Result, error) {
	if generation < 1 || now.IsZero() {
		return Result{}, fmt.Errorf("invalid backup generation or timestamp")
	}
	paths, err := tenantfs.Derive(dataRoot, tenantID)
	if err != nil {
		return Result{}, err
	}
	if err := checkpointAndVerify(paths.Database); err != nil {
		return Result{}, err
	}
	files, err := sourceFiles(paths)
	if err != nil {
		return Result{}, err
	}
	backupRoot, err = filepath.Abs(strings.TrimSpace(backupRoot))
	if err != nil || strings.TrimSpace(backupRoot) == "" {
		return Result{}, fmt.Errorf("invalid backup root")
	}
	if err := ensureRealDirectory(backupRoot); err != nil {
		return Result{}, err
	}
	tenantBackupRoot := filepath.Join(backupRoot, filepath.Base(paths.Root))
	if err := ensureRealDirectory(tenantBackupRoot); err != nil {
		return Result{}, err
	}
	nonce, err := randomHex(8)
	if err != nil {
		return Result{}, err
	}
	name := fmt.Sprintf("g%020d-%s-%s", generation, now.UTC().Format("20060102T150405Z"), nonce)
	finalPath := filepath.Join(tenantBackupRoot, name)
	stagingPath := finalPath + ".tmp"
	if err := os.Mkdir(stagingPath, 0o700); err != nil {
		return Result{}, err
	}
	committed := false
	defer func() {
		if !committed {
			_ = os.RemoveAll(stagingPath)
		}
	}()
	manifest := Manifest{
		Version: manifestVersion, TenantID: tenantID, PlacementGeneration: generation,
		CreatedAt: now.UTC().Format(time.RFC3339Nano), Files: make([]File, 0, len(files)),
	}
	for _, item := range files {
		copied, err := copyFile(ctx, item.absolute, filepath.Join(stagingPath, filepath.FromSlash(item.relative)))
		if err != nil {
			return Result{}, err
		}
		copied.Path = item.relative
		manifest.Files = append(manifest.Files, copied)
	}
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		return Result{}, err
	}
	manifestPath := filepath.Join(stagingPath, "manifest.json")
	manifestFile, err := os.OpenFile(manifestPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return Result{}, err
	}
	if _, err := manifestFile.Write(manifestBytes); err != nil {
		_ = manifestFile.Close()
		return Result{}, err
	}
	if err := manifestFile.Sync(); err != nil {
		_ = manifestFile.Close()
		return Result{}, err
	}
	if err := manifestFile.Close(); err != nil {
		return Result{}, err
	}
	if err := os.Rename(stagingPath, finalPath); err != nil {
		return Result{}, err
	}
	committed = true
	digest := sha256.Sum256(manifestBytes)
	backupRef := filepath.ToSlash(filepath.Join(filepath.Base(paths.Root), name))
	if !validBackupRef(backupRef) {
		return Result{}, fmt.Errorf("generated backup reference is invalid")
	}
	return Result{Path: finalPath, BackupRef: backupRef, Manifest: manifest, ManifestSHA256: hex.EncodeToString(digest[:])}, nil
}

func validManifestPath(path string) bool {
	if path == "relay.db" {
		return true
	}
	parts := strings.Split(path, "/")
	return len(parts) == 2 && parts[0] == "attachments" && validAttachmentName(parts[1])
}

func Verify(ctx context.Context, backupPath, tenantID string, generation int64, expectedManifestSHA256 string) (Manifest, error) {
	backupPath, err := filepath.Abs(strings.TrimSpace(backupPath))
	if err != nil || strings.TrimSpace(backupPath) == "" {
		return Manifest{}, fmt.Errorf("invalid backup path")
	}
	info, err := os.Lstat(backupPath)
	if err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return Manifest{}, fmt.Errorf("backup path is not a real directory")
	}
	manifestPath := filepath.Join(backupPath, "manifest.json")
	manifestInfo, err := regularFile(manifestPath)
	if err != nil || manifestInfo.Size() > 1<<20 {
		return Manifest{}, fmt.Errorf("invalid backup manifest")
	}
	manifestBytes, err := os.ReadFile(manifestPath)
	if err != nil {
		return Manifest{}, err
	}
	digest := sha256.Sum256(manifestBytes)
	if expectedManifestSHA256 != hex.EncodeToString(digest[:]) {
		return Manifest{}, fmt.Errorf("backup manifest digest mismatch")
	}
	decoder := json.NewDecoder(bytes.NewReader(manifestBytes))
	decoder.DisallowUnknownFields()
	var manifest Manifest
	if err := decoder.Decode(&manifest); err != nil || decoder.Decode(&struct{}{}) != io.EOF {
		return Manifest{}, fmt.Errorf("invalid backup manifest JSON")
	}
	if manifest.Version != manifestVersion || manifest.TenantID != tenantID || manifest.PlacementGeneration != generation || len(manifest.Files) == 0 {
		return Manifest{}, fmt.Errorf("backup manifest fence mismatch")
	}
	expectedFiles := map[string]struct{}{"manifest.json": {}}
	previous := ""
	for _, item := range manifest.Files {
		if !validManifestPath(item.Path) || item.Path <= previous || item.Size < 0 || len(item.SHA256) != 64 {
			return Manifest{}, fmt.Errorf("invalid backup file manifest")
		}
		previous = item.Path
		absolute := filepath.Join(backupPath, filepath.FromSlash(item.Path))
		if !strings.HasPrefix(absolute, backupPath+string(os.PathSeparator)) {
			return Manifest{}, fmt.Errorf("backup file escaped root")
		}
		size, checksum, err := hashFile(ctx, absolute)
		if err != nil || size != item.Size || checksum != item.SHA256 {
			return Manifest{}, fmt.Errorf("backup file verification failed for %s", item.Path)
		}
		expectedFiles[filepath.Clean(item.Path)] = struct{}{}
	}
	err = filepath.WalkDir(backupPath, func(path string, entry os.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if path == backupPath {
			return nil
		}
		relative, err := filepath.Rel(backupPath, path)
		if err != nil {
			return err
		}
		if entry.Type()&os.ModeSymlink != 0 {
			return fmt.Errorf("backup contains a symbolic link")
		}
		if entry.IsDir() {
			if filepath.Clean(relative) != "attachments" {
				return fmt.Errorf("backup contains an unexpected directory")
			}
			return nil
		}
		if _, ok := expectedFiles[filepath.Clean(relative)]; !ok {
			return fmt.Errorf("backup contains an unexpected file %q", relative)
		}
		return nil
	})
	if err != nil {
		return Manifest{}, err
	}
	if err := verifySQLiteReadOnly(filepath.Join(backupPath, "relay.db")); err != nil {
		return Manifest{}, err
	}
	return manifest, nil
}

func Restore(ctx context.Context, backupPath, dataRoot, tenantID string, generation int64, manifestSHA256 string) (tenantfs.Paths, error) {
	manifest, err := Verify(ctx, backupPath, tenantID, generation, manifestSHA256)
	if err != nil {
		return tenantfs.Paths{}, err
	}
	paths, err := tenantfs.Derive(dataRoot, tenantID)
	if err != nil {
		return tenantfs.Paths{}, err
	}
	dataRoot = filepath.Dir(filepath.Dir(paths.Root))
	if err := ensureRealDirectory(dataRoot); err != nil {
		return tenantfs.Paths{}, err
	}
	tenantsRoot := filepath.Dir(paths.Root)
	if err := ensureRealDirectory(tenantsRoot); err != nil {
		return tenantfs.Paths{}, err
	}
	if info, err := os.Lstat(paths.Root); err == nil {
		if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
			return tenantfs.Paths{}, fmt.Errorf("target tenant path is not a real directory")
		}
		receiptBytes, readErr := os.ReadFile(filepath.Join(paths.Root, ".restore-receipt.json"))
		var receipt restoreReceipt
		decoder := json.NewDecoder(bytes.NewReader(receiptBytes))
		decoder.DisallowUnknownFields()
		if readErr == nil && decoder.Decode(&receipt) == nil && decoder.Decode(&struct{}{}) == io.EOF &&
			receipt.Version == manifestVersion && receipt.TenantID == tenantID &&
			receipt.PlacementGeneration == generation && receipt.ManifestSHA256 == manifestSHA256 {
			return paths, nil
		}
		return tenantfs.Paths{}, fmt.Errorf("target tenant directory already exists with another restore fence")
	} else if !os.IsNotExist(err) {
		return tenantfs.Paths{}, err
	}
	nonce, err := randomHex(8)
	if err != nil {
		return tenantfs.Paths{}, err
	}
	staging := filepath.Join(tenantsRoot, ".restore-"+filepath.Base(paths.Root)+"-"+nonce)
	if err := os.Mkdir(staging, 0o700); err != nil {
		return tenantfs.Paths{}, err
	}
	committed := false
	defer func() {
		if !committed {
			_ = os.RemoveAll(staging)
		}
	}()
	for _, item := range manifest.Files {
		if _, err := copyFile(ctx, filepath.Join(backupPath, filepath.FromSlash(item.Path)), filepath.Join(staging, filepath.FromSlash(item.Path))); err != nil {
			return tenantfs.Paths{}, err
		}
	}
	if err := verifySQLiteReadOnly(filepath.Join(staging, "relay.db")); err != nil {
		return tenantfs.Paths{}, err
	}
	receiptBytes, err := json.Marshal(restoreReceipt{
		Version: manifestVersion, TenantID: tenantID,
		PlacementGeneration: generation, ManifestSHA256: manifestSHA256,
	})
	if err != nil {
		return tenantfs.Paths{}, err
	}
	receipt, err := os.OpenFile(filepath.Join(staging, ".restore-receipt.json"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return tenantfs.Paths{}, err
	}
	if _, err := receipt.Write(receiptBytes); err != nil {
		_ = receipt.Close()
		return tenantfs.Paths{}, err
	}
	if err := receipt.Sync(); err != nil {
		_ = receipt.Close()
		return tenantfs.Paths{}, err
	}
	if err := receipt.Close(); err != nil {
		return tenantfs.Paths{}, err
	}
	if err := os.Rename(staging, paths.Root); err != nil {
		return tenantfs.Paths{}, err
	}
	committed = true
	return paths, nil
}