banger/internal/daemon/workspace.go

package daemon

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"banger/internal/api"
	"banger/internal/model"
	"banger/internal/system"
)

const workspaceShallowFetchDepth = 10

type workspaceRepoSpec struct {
	SourcePath    string
	RepoRoot      string
	RepoName      string
	HeadCommit    string
	CurrentBranch string
	BranchName    string
	BaseCommit    string
	OriginURL     string
	GitUserName   string
	GitUserEmail  string
	OverlayPaths  []string
	Submodules    []string
}

func (d *Daemon) PrepareVMWorkspace(ctx context.Context, params api.VMWorkspacePrepareParams) (model.WorkspacePrepareResult, error) {
	mode, err := parseWorkspacePrepareMode(params.Mode)
	if err != nil {
		return model.WorkspacePrepareResult{}, err
	}
	guestPath := strings.TrimSpace(params.GuestPath)
	if guestPath == "" {
		guestPath = "/root/repo"
	}
	branchName := strings.TrimSpace(params.Branch)
	fromRef := strings.TrimSpace(params.From)
	if branchName != "" && fromRef == "" {
		fromRef = "HEAD"
	}
	if branchName == "" && strings.TrimSpace(params.From) != "" {
		return model.WorkspacePrepareResult{}, errors.New("workspace from requires branch")
	}
	var prepared model.WorkspacePrepareResult
	_, err = d.withVMLockByRef(ctx, params.IDOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		if vm.State != model.VMStateRunning || !system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
			return model.VMRecord{}, fmt.Errorf("vm %q is not running", vm.Name)
		}
		result, err := d.prepareVMWorkspaceLocked(ctx, vm, strings.TrimSpace(params.SourcePath), guestPath, branchName, fromRef, mode, params.ReadOnly)
		if err != nil {
			return model.VMRecord{}, err
		}
		prepared = result
		return vm, nil
	})
	return prepared, err
}

func (d *Daemon) prepareVMWorkspaceLocked(ctx context.Context, vm model.VMRecord, sourcePath, guestPath, branchName, fromRef string, mode model.WorkspacePrepareMode, readOnly bool) (model.WorkspacePrepareResult, error) {
	spec, err := inspectWorkspaceRepo(ctx, sourcePath, branchName, fromRef)
	if err != nil {
		return model.WorkspacePrepareResult{}, err
	}
	if len(spec.Submodules) > 0 && mode != model.WorkspacePrepareModeFullCopy {
		return model.WorkspacePrepareResult{}, fmt.Errorf("workspace mode %q does not support git submodules in %s (%s); use --mode full_copy", mode, spec.RepoRoot, strings.Join(spec.Submodules, ", "))
	}
	address := net.JoinHostPort(vm.Runtime.GuestIP, "22")
	if err := d.waitForGuestSSH(ctx, address, 250*time.Millisecond); err != nil {
		return model.WorkspacePrepareResult{}, fmt.Errorf("guest ssh unavailable: %w", err)
	}
	client, err := d.dialGuest(ctx, address)
	if err != nil {
		return model.WorkspacePrepareResult{}, fmt.Errorf("dial guest ssh: %w", err)
	}
	defer client.Close()
	if err := importWorkspaceRepoToGuest(ctx, client, spec, guestPath, mode); err != nil {
		return model.WorkspacePrepareResult{}, err
	}
	if readOnly {
		var chmodLog bytes.Buffer
		chmodScript := fmt.Sprintf("set -euo pipefail\nchmod -R a-w %s\n", guestShellQuote(guestPath))
		if err := client.RunScript(ctx, chmodScript, &chmodLog); err != nil {
			return model.WorkspacePrepareResult{}, formatGuestSessionStepError("set workspace readonly", err, chmodLog.String())
		}
	}
	return model.WorkspacePrepareResult{
		VMID:          vm.ID,
		SourcePath:    spec.SourcePath,
		RepoRoot:      spec.RepoRoot,
		RepoName:      spec.RepoName,
		GuestPath:     guestPath,
		Mode:          mode,
		ReadOnly:      readOnly,
		HeadCommit:    spec.HeadCommit,
		CurrentBranch: spec.CurrentBranch,
		BranchName:    spec.BranchName,
		BaseCommit:    spec.BaseCommit,
		PreparedAt:    model.Now(),
	}, nil
}

func inspectWorkspaceRepo(ctx context.Context, rawPath, branchName, fromRef string) (workspaceRepoSpec, error) {
	sourcePath, err := resolveWorkspaceSourcePath(rawPath)
	if err != nil {
		return workspaceRepoSpec{}, err
	}
	repoRoot, err := workspaceGitTrimmedOutput(ctx, sourcePath, "rev-parse", "--show-toplevel")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("%s is not inside a git repository", sourcePath)
	}
	isBare, err := workspaceGitTrimmedOutput(ctx, repoRoot, "rev-parse", "--is-bare-repository")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("inspect git repository %s: %w", repoRoot, err)
	}
	if isBare == "true" {
		return workspaceRepoSpec{}, fmt.Errorf("workspace prepare requires a non-bare git repository: %s", repoRoot)
	}
	submodules, err := listWorkspaceSubmodules(ctx, repoRoot)
	if err != nil {
		return workspaceRepoSpec{}, err
	}
	headCommit, err := workspaceGitTrimmedOutput(ctx, repoRoot, "rev-parse", "HEAD^{commit}")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("git repository %s must have at least one commit", repoRoot)
	}
	currentBranch, err := workspaceGitTrimmedOutput(ctx, repoRoot, "branch", "--show-current")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("resolve current branch for %s: %w", repoRoot, err)
	}
	baseCommit := headCommit
	branchName = strings.TrimSpace(branchName)
	if branchName != "" {
		baseCommit, err = workspaceGitTrimmedOutput(ctx, repoRoot, "rev-parse", fromRef+"^{commit}")
		if err != nil {
			return workspaceRepoSpec{}, fmt.Errorf("resolve workspace from %q: %w", fromRef, err)
		}
	}
	gitUserName, err := workspaceGitResolvedConfigValue(ctx, repoRoot, "user.name")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("resolve git user.name for %s: %w", repoRoot, err)
	}
	gitUserEmail, err := workspaceGitResolvedConfigValue(ctx, repoRoot, "user.email")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("resolve git user.email for %s: %w", repoRoot, err)
	}
	originURL, err := workspaceGitResolvedConfigValue(ctx, repoRoot, "remote.origin.url")
	if err != nil {
		return workspaceRepoSpec{}, fmt.Errorf("resolve origin url for %s: %w", repoRoot, err)
	}
	overlayPaths, err := listWorkspaceOverlayPaths(ctx, repoRoot)
	if err != nil {
		return workspaceRepoSpec{}, err
	}
	return workspaceRepoSpec{
		SourcePath:    sourcePath,
		RepoRoot:      repoRoot,
		RepoName:      filepath.Base(repoRoot),
		HeadCommit:    headCommit,
		CurrentBranch: currentBranch,
		BranchName:    branchName,
		BaseCommit:    baseCommit,
		OriginURL:     originURL,
		GitUserName:   gitUserName,
		GitUserEmail:  gitUserEmail,
		OverlayPaths:  overlayPaths,
		Submodules:    submodules,
	}, nil
}

func importWorkspaceRepoToGuest(ctx context.Context, client guestSSHClient, spec workspaceRepoSpec, guestPath string, mode model.WorkspacePrepareMode) error {
	switch mode {
	case model.WorkspacePrepareModeFullCopy:
		var copyLog bytes.Buffer
		command := fmt.Sprintf("rm -rf %s && mkdir -p %s && tar -o -C %s --strip-components=1 -xf -", guestShellQuote(guestPath), guestShellQuote(guestPath), guestShellQuote(guestPath))
		if err := client.StreamTar(ctx, spec.RepoRoot, command, &copyLog); err != nil {
			return formatGuestSessionStepError("copy full workspace", err, copyLog.String())
		}
		var finalizeLog bytes.Buffer
		if err := client.RunScript(ctx, workspaceFinalizeScript(spec, guestPath, mode), &finalizeLog); err != nil {
			return formatGuestSessionStepError("finalize full workspace", err, finalizeLog.String())
		}
		return nil
	case model.WorkspacePrepareModeMetadataOnly, model.WorkspacePrepareModeShallowOverlay:
		repoCopyDir, cleanup, err := prepareWorkspaceRepoCopy(ctx, spec)
		if err != nil {
			return err
		}
		defer cleanup()
		var copyLog bytes.Buffer
		command := fmt.Sprintf("rm -rf %s && mkdir -p %s && tar -o -C %s --strip-components=1 -xf -", guestShellQuote(guestPath), guestShellQuote(guestPath), guestShellQuote(guestPath))
		if err := client.StreamTar(ctx, repoCopyDir, command, &copyLog); err != nil {
			return formatGuestSessionStepError("copy guest git metadata", err, copyLog.String())
		}
		var scriptLog bytes.Buffer
		if err := client.RunScript(ctx, workspaceFinalizeScript(spec, guestPath, mode), &scriptLog); err != nil {
			return formatGuestSessionStepError("prepare guest checkout", err, scriptLog.String())
		}
		if mode == model.WorkspacePrepareModeMetadataOnly {
			return nil
		}
		var overlayLog bytes.Buffer
		command = fmt.Sprintf("tar -o -C %s --strip-components=1 -xf -", guestShellQuote(guestPath))
		if err := client.StreamTarEntries(ctx, spec.RepoRoot, spec.OverlayPaths, command, &overlayLog); err != nil {
			return formatGuestSessionStepError("overlay workspace working tree", err, overlayLog.String())
		}
		return nil
	default:
		return fmt.Errorf("unsupported workspace mode %q", mode)
	}
}

func workspaceFinalizeScript(spec workspaceRepoSpec, guestPath string, mode model.WorkspacePrepareMode) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	fmt.Fprintf(&script, "DIR=%s\n", guestShellQuote(guestPath))
	script.WriteString("git config --global --add safe.directory \"$DIR\"\n")
	if mode != model.WorkspacePrepareModeFullCopy {
		script.WriteString("find \"$DIR\" -mindepth 1 -maxdepth 1 ! -name .git -exec rm -rf {} +\n")
	}
	switch {
	case strings.TrimSpace(spec.BranchName) != "":
		fmt.Fprintf(&script, "git -C \"$DIR\" checkout -B %s %s\n", guestShellQuote(spec.BranchName), guestShellQuote(spec.BaseCommit))
	case strings.TrimSpace(spec.CurrentBranch) != "":
		fmt.Fprintf(&script, "git -C \"$DIR\" checkout -B %s %s\n", guestShellQuote(spec.CurrentBranch), guestShellQuote(spec.HeadCommit))
	default:
		fmt.Fprintf(&script, "git -C \"$DIR\" checkout --detach %s\n", guestShellQuote(spec.HeadCommit))
	}
	if strings.TrimSpace(spec.GitUserName) != "" && strings.TrimSpace(spec.GitUserEmail) != "" {
		fmt.Fprintf(&script, "git -C \"$DIR\" config user.name %s\n", guestShellQuote(spec.GitUserName))
		fmt.Fprintf(&script, "git -C \"$DIR\" config user.email %s\n", guestShellQuote(spec.GitUserEmail))
	}
	return script.String()
}

func prepareWorkspaceRepoCopy(ctx context.Context, spec workspaceRepoSpec) (string, func(), error) {
	tempRoot, err := os.MkdirTemp("", "banger-workspace-*")
	if err != nil {
		return "", nil, err
	}
	cleanup := func() { _ = os.RemoveAll(tempRoot) }
	repoCopyDir := filepath.Join(tempRoot, spec.RepoName)
	cloneArgs := []string{"clone", "--no-checkout", "--depth", fmt.Sprintf("%d", workspaceShallowFetchDepth)}
	if strings.TrimSpace(spec.CurrentBranch) != "" {
		cloneArgs = append(cloneArgs, "--single-branch", "--branch", spec.CurrentBranch)
	}
	cloneArgs = append(cloneArgs, workspaceGitFileURL(spec.RepoRoot), repoCopyDir)
	if err := workspaceRunHostCommand(ctx, "git", cloneArgs...); err != nil {
		cleanup()
		return "", nil, fmt.Errorf("clone shallow workspace repo copy: %w", err)
	}
	checkoutCommit := spec.HeadCommit
	if strings.TrimSpace(spec.BranchName) != "" {
		checkoutCommit = spec.BaseCommit
	}
	if err := workspaceRunHostCommand(ctx, "git", "-C", repoCopyDir, "cat-file", "-e", checkoutCommit+"^{commit}"); err != nil {
		if err := workspaceRunHostCommand(ctx, "git", "-C", repoCopyDir, "fetch", "--depth", fmt.Sprintf("%d", workspaceShallowFetchDepth), workspaceGitFileURL(spec.RepoRoot), checkoutCommit); err != nil {
			cleanup()
			return "", nil, fmt.Errorf("fetch shallow workspace repo commit %s: %w", checkoutCommit, err)
		}
	}
	if strings.TrimSpace(spec.OriginURL) != "" {
		if err := workspaceRunHostCommand(ctx, "git", "-C", repoCopyDir, "remote", "set-url", "origin", spec.OriginURL); err != nil {
			cleanup()
			return "", nil, fmt.Errorf("set workspace origin remote: %w", err)
		}
	} else {
		if err := workspaceRunHostCommand(ctx, "git", "-C", repoCopyDir, "remote", "remove", "origin"); err != nil {
			cleanup()
			return "", nil, fmt.Errorf("remove workspace placeholder origin remote: %w", err)
		}
	}
	return repoCopyDir, cleanup, nil
}

func resolveWorkspaceSourcePath(rawPath string) (string, error) {
	if strings.TrimSpace(rawPath) == "" {
		return "", errors.New("workspace source path is required")
	}
	absPath, err := filepath.Abs(rawPath)
	if err != nil {
		return "", err
	}
	info, err := os.Stat(absPath)
	if err != nil {
		return "", err
	}
	if !info.IsDir() {
		return "", fmt.Errorf("%s is not a directory", absPath)
	}
	return absPath, nil
}

func listWorkspaceSubmodules(ctx context.Context, repoRoot string) ([]string, error) {
	output, err := workspaceGitOutput(ctx, repoRoot, "ls-files", "--stage", "-z")
	if err != nil {
		return nil, fmt.Errorf("inspect workspace git index for %s: %w", repoRoot, err)
	}
	var submodules []string
	for _, record := range workspaceParseNullSeparatedOutput(output) {
		if !strings.HasPrefix(record, "160000 ") {
			continue
		}
		_, path, ok := strings.Cut(record, "	")
		if !ok {
			continue
		}
		submodules = append(submodules, strings.TrimSpace(path))
	}
	sort.Strings(submodules)
	return submodules, nil
}

func listWorkspaceOverlayPaths(ctx context.Context, repoRoot string) ([]string, error) {
	trackedOutput, err := workspaceGitOutput(ctx, repoRoot, "ls-files", "-z")
	if err != nil {
		return nil, fmt.Errorf("list tracked files for %s: %w", repoRoot, err)
	}
	untrackedOutput, err := workspaceGitOutput(ctx, repoRoot, "ls-files", "--others", "--exclude-standard", "-z")
	if err != nil {
		return nil, fmt.Errorf("list untracked files for %s: %w", repoRoot, err)
	}
	paths := make([]string, 0)
	seen := make(map[string]struct{})
	for _, relPath := range workspaceParseNullSeparatedOutput(trackedOutput) {
		if relPath == "" {
			continue
		}
		if _, err := os.Lstat(filepath.Join(repoRoot, relPath)); err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		seen[relPath] = struct{}{}
		paths = append(paths, relPath)
	}
	for _, relPath := range workspaceParseNullSeparatedOutput(untrackedOutput) {
		if relPath == "" {
			continue
		}
		if _, ok := seen[relPath]; ok {
			continue
		}
		seen[relPath] = struct{}{}
		paths = append(paths, relPath)
	}
	sort.Strings(paths)
	return paths, nil
}

func parseWorkspacePrepareMode(raw string) (model.WorkspacePrepareMode, error) {
	switch strings.TrimSpace(raw) {
	case "", string(model.WorkspacePrepareModeShallowOverlay):
		return model.WorkspacePrepareModeShallowOverlay, nil
	case string(model.WorkspacePrepareModeFullCopy):
		return model.WorkspacePrepareModeFullCopy, nil
	case string(model.WorkspacePrepareModeMetadataOnly):
		return model.WorkspacePrepareModeMetadataOnly, nil
	default:
		return "", fmt.Errorf("unsupported workspace mode %q", raw)
	}
}

func workspaceGitOutput(ctx context.Context, dir string, args ...string) ([]byte, error) {
	fullArgs := make([]string, 0, len(args)+2)
	if strings.TrimSpace(dir) != "" {
		fullArgs = append(fullArgs, "-C", dir)
	}
	fullArgs = append(fullArgs, args...)
	return guestSessionHostCommandOutputFunc(ctx, "git", fullArgs...)
}

func workspaceGitTrimmedOutput(ctx context.Context, dir string, args ...string) (string, error) {
	output, err := workspaceGitOutput(ctx, dir, args...)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(output)), nil
}

func workspaceGitResolvedConfigValue(ctx context.Context, dir, key string) (string, error) {
	return workspaceGitTrimmedOutput(ctx, dir, "config", "--default", "", "--get", key)
}

func workspaceParseNullSeparatedOutput(output []byte) []string {
	chunks := bytes.Split(output, []byte{0})
	values := make([]string, 0, len(chunks))
	for _, chunk := range chunks {
		value := strings.TrimSpace(string(chunk))
		if value == "" {
			continue
		}
		values = append(values, value)
	}
	return values
}

func workspaceRunHostCommand(ctx context.Context, name string, args ...string) error {
	_, err := guestSessionHostCommandOutputFunc(ctx, name, args...)
	return err
}

func workspaceGitFileURL(path string) string {
	return (&url.URL{Scheme: "file", Path: filepath.ToSlash(path)}).String()
}