banger/internal/daemon/workspace.go

package daemon

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"

	"banger/internal/api"
	ws "banger/internal/daemon/workspace"
	"banger/internal/model"
)

// workspaceInspectRepoHook + workspaceImportHook dispatch through the
// per-instance Daemon seams when set, falling back to the real
// workspace package implementations. Keeping the fallbacks here (as
// opposed to always requiring callers to populate s.workspaceInspectRepo
// in a constructor) lets tests selectively override one hook without
// having to wire both.
func (s *WorkspaceService) workspaceInspectRepoHook(ctx context.Context, sourcePath, branchName, fromRef string, includeUntracked bool) (ws.RepoSpec, error) {
	if s != nil && s.workspaceInspectRepo != nil {
		return s.workspaceInspectRepo(ctx, sourcePath, branchName, fromRef, includeUntracked)
	}
	return s.inspector().InspectRepo(ctx, sourcePath, branchName, fromRef, includeUntracked)
}

func (s *WorkspaceService) workspaceImportHook(ctx context.Context, client ws.GuestClient, spec ws.RepoSpec, guestPath string, mode model.WorkspacePrepareMode) error {
	if s != nil && s.workspaceImport != nil {
		return s.workspaceImport(ctx, client, spec, guestPath, mode)
	}
	return s.inspector().ImportRepoToGuest(ctx, client, spec, guestPath, mode)
}

// inspector returns the service's workspace Inspector, falling back to
// a fresh real-runner Inspector when callers constructed the service
// without wiring one. Keeping the fallback here lets test literals
// that don't care about the Inspector still function without a manual
// NewInspector() call.
func (s *WorkspaceService) inspector() *ws.Inspector {
	if s != nil && s.repoInspector != nil {
		return s.repoInspector
	}
	return ws.NewInspector()
}

func (s *WorkspaceService) ExportVMWorkspace(ctx context.Context, params api.WorkspaceExportParams) (api.WorkspaceExportResult, error) {
	guestPath := strings.TrimSpace(params.GuestPath)
	if guestPath == "" {
		guestPath = "/root/repo"
	}
	vm, err := s.vmResolver(ctx, params.IDOrName)
	if err != nil {
		return api.WorkspaceExportResult{}, err
	}
	if !s.aliveChecker(vm) {
		return api.WorkspaceExportResult{}, fmt.Errorf("vm %q is not running", vm.Name)
	}
	// Serialise with any in-flight workspace.prepare on the same VM so
	// we never snapshot a half-streamed tar. Does not block vm stop /
	// delete / restart — those only take the VM mutex.
	unlock := s.workspaceLocks.lock(vm.ID)
	defer unlock()

	client, err := s.dialGuest(ctx, net.JoinHostPort(vm.Runtime.GuestIP, "22"))
	if err != nil {
		return api.WorkspaceExportResult{}, fmt.Errorf("dial guest: %w", err)
	}
	defer client.Close()

	// diffRef is the git ref everything is diffed against.
	// When the caller supplies BaseCommit (the HEAD at workspace.prepare time),
	// we diff against that fixed point so committed guest changes are included.
	// Without it we fall back to HEAD, which silently drops them.
	diffRef := strings.TrimSpace(params.BaseCommit)
	if diffRef == "" {
		diffRef = "HEAD"
	}

	// Both scripts run `git add -A` to capture the working tree
	// (committed deltas + uncommitted modifications + untracked files
	// minus .gitignore), but they route it through a throwaway index
	// file instead of .git/index. Export is an observation step; the
	// user's real staging area must stay exactly as they left it.
	patchScript := exportScript(guestPath, diffRef, "--binary")
	patch, err := client.RunScriptOutput(ctx, patchScript)
	if err != nil {
		return api.WorkspaceExportResult{}, fmt.Errorf("export workspace diff: %w", err)
	}

	namesScript := exportScript(guestPath, diffRef, "--name-only")
	namesOut, _ := client.RunScriptOutput(ctx, namesScript)
	var changed []string
	for _, line := range strings.Split(strings.TrimSpace(string(namesOut)), "\n") {
		if line = strings.TrimSpace(line); line != "" {
			changed = append(changed, line)
		}
	}

	return api.WorkspaceExportResult{
		GuestPath:    guestPath,
		BaseCommit:   diffRef,
		Patch:        patch,
		ChangedFiles: changed,
		HasChanges:   len(patch) > 0,
	}, nil
}

// exportScript emits a shell snippet that diffs the working tree at
// guestPath against diffRef (HEAD or a commit SHA) WITHOUT touching
// the repo's real index. diffFlag selects the git-diff output mode
// ("--binary" for the patch body, "--name-only" for the file list).
//
// Mechanics: seed a temp index from diffRef's tree via git read-tree,
// restage the working tree into that temp index with GIT_INDEX_FILE,
// then emit the diff. The temp index is rm'd on exit via trap.
func exportScript(guestPath, diffRef, diffFlag string) string {
	return fmt.Sprintf(
		"set -euo pipefail\n"+
			"cd %s\n"+
			"tmp_idx=\"$(mktemp \"${TMPDIR:-/tmp}/banger-export.XXXXXX\")\"\n"+
			"trap 'rm -f \"$tmp_idx\"' EXIT\n"+
			"git read-tree %s --index-output=\"$tmp_idx\"\n"+
			"GIT_INDEX_FILE=\"$tmp_idx\" git add -A\n"+
			"GIT_INDEX_FILE=\"$tmp_idx\" git diff --cached %s %s\n",
		ws.ShellQuote(guestPath),
		ws.ShellQuote(diffRef),
		ws.ShellQuote(diffRef),
		diffFlag,
	)
}

func (s *WorkspaceService) PrepareVMWorkspace(ctx context.Context, params api.VMWorkspacePrepareParams) (model.WorkspacePrepareResult, error) {
	mode, err := ws.ParsePrepareMode(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")
	}

	// Phase 1: acquire the VM mutex ONLY long enough to verify state
	// and snapshot the fields we need (IP, PID, api sock). Release it
	// before any SSH or tar I/O so this slow operation cannot block
	// vm stop / vm delete / vm restart on the same VM.
	vm, err := s.withVMLockByRef(ctx, params.IDOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		if !s.aliveChecker(vm) {
			return model.VMRecord{}, fmt.Errorf("vm %q is not running", vm.Name)
		}
		return vm, nil
	})
	if err != nil {
		return model.WorkspacePrepareResult{}, err
	}

	// Phase 2: serialise concurrent workspace operations on THIS vm
	// (so two prepares don't interleave tar streams), but do not
	// block lifecycle ops. If the VM gets stopped or deleted mid-
	// flight, the SSH dial or stream will fail naturally; ctx
	// cancellation propagates through.
	unlock := s.workspaceLocks.lock(vm.ID)
	defer unlock()

	return s.prepareVMWorkspaceGuestIO(ctx, vm, strings.TrimSpace(params.SourcePath), guestPath, branchName, fromRef, mode, params.ReadOnly, params.IncludeUntracked)
}

// prepareVMWorkspaceGuestIO performs the actual guest-side work:
// inspect the local repo, dial SSH, stream the tar, optionally chmod
// readonly. It is called without holding the VM mutex.
func (s *WorkspaceService) prepareVMWorkspaceGuestIO(ctx context.Context, vm model.VMRecord, sourcePath, guestPath, branchName, fromRef string, mode model.WorkspacePrepareMode, readOnly, includeUntracked bool) (model.WorkspacePrepareResult, error) {
	spec, err := s.workspaceInspectRepoHook(ctx, sourcePath, branchName, fromRef, includeUntracked)
	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 := s.waitGuestSSH(ctx, address, 250*time.Millisecond); err != nil {
		return model.WorkspacePrepareResult{}, fmt.Errorf("guest ssh unavailable: %w", err)
	}
	client, err := s.dialGuest(ctx, address)
	if err != nil {
		return model.WorkspacePrepareResult{}, fmt.Errorf("dial guest ssh: %w", err)
	}
	defer client.Close()
	if err := s.workspaceImportHook(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", ws.ShellQuote(guestPath))
		if err := client.RunScript(ctx, chmodScript, &chmodLog); err != nil {
			return model.WorkspacePrepareResult{}, ws.FormatStepError("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
}