Agent-Squad/squad/deliverable.go

package squad

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

func slashClean(p string) string {
	return filepath.ToSlash(filepath.Clean(p))
}

// MatchDeliverableRel reports whether a runner rel_path refers to target
// inside workDir. Accepts output.md, ./output.md, absolute paths under workDir,
// and a duplicated step prefix (tasks/01_draft/output.md when workDir already
// ends with tasks/01_draft). Paths that escape workDir never match.
func MatchDeliverableRel(workDir, relPath, target string) bool {
	target = slashClean(strings.TrimSpace(target))
	if target == "" || target == "." || strings.HasPrefix(target, "../") || strings.Contains(target, "..") {
		return false
	}
	raw := strings.TrimSpace(relPath)
	if raw == "" {
		return false
	}
	abs := raw
	if !filepath.IsAbs(raw) {
		abs = filepath.Join(workDir, raw)
	}
	abs = filepath.Clean(abs)
	rel, err := filepath.Rel(filepath.Clean(workDir), abs)
	if err != nil {
		return false
	}
	rel = slashClean(rel)
	if rel == ".." || strings.HasPrefix(rel, "../") {
		return false
	}
	if rel == target {
		return true
	}
	if filepath.Base(rel) != filepath.Base(target) {
		return false
	}
	prefix := slashClean(filepath.Dir(rel))
	wd := slashClean(workDir)
	return prefix != "." && (strings.HasSuffix(wd, "/"+prefix) || strings.HasSuffix(wd, prefix))
}

// CollectTargetPatches maps incoming runner keys onto canonical target names
// and also accepts a non-empty file already at workDir/target. Empty bodies
// do not count — callers must not invent a deliverable.
func CollectTargetPatches(workDir string, targets []string, incoming map[string][]byte) map[string][]byte {
	out := make(map[string][]byte, len(targets))
	for _, t := range targets {
		t = strings.TrimSpace(t)
		if t == "" {
			continue
		}
		for k, v := range incoming {
			if len(bytes.TrimSpace(v)) == 0 {
				continue
			}
			if MatchDeliverableRel(workDir, k, t) {
				out[t] = v
				break
			}
		}
		if _, ok := out[t]; ok {
			continue
		}
		p := filepath.Join(workDir, t)
		if b, err := os.ReadFile(p); err == nil && len(bytes.TrimSpace(b)) > 0 {
			out[t] = b
		}
	}
	return out
}

// MaterializeTargets writes collected non-empty target bodies to workDir/<target>.
// Returns a missing-target error without writing a stub if any required file is absent.
func MaterializeTargets(workDir string, targets []string, incoming map[string][]byte) (map[string][]byte, error) {
	collected := CollectTargetPatches(workDir, targets, incoming)
	for _, t := range targets {
		t = strings.TrimSpace(t)
		if t == "" {
			continue
		}
		body, ok := collected[t]
		if !ok || len(bytes.TrimSpace(body)) == 0 {
			return collected, fmt.Errorf("missing target deliverable %s", t)
		}
		dest := filepath.Join(workDir, filepath.FromSlash(t))
		if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
			return collected, err
		}
		if err := rejectSymlinkEscape(workDir, dest); err != nil {
			return collected, err
		}
		// Container-written files may be root-owned; make writable before overwrite.
		_ = os.Chmod(dest, 0o666)
		if err := os.WriteFile(dest, body, 0o644); err != nil {
			return collected, err
		}
	}
	return collected, nil
}

func rejectSymlinkEscape(workDir, dest string) error {
	root, err := filepath.Abs(workDir)
	if err != nil {
		return err
	}
	absDest, err := filepath.Abs(dest)
	if err != nil {
		return err
	}
	rel, err := filepath.Rel(root, absDest)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return fmt.Errorf("typed result path rejected")
	}
	resolvedRoot, err := filepath.EvalSymlinks(root)
	if err != nil {
		resolvedRoot = root
	}
	parent := filepath.Dir(absDest)
	resolvedParent, err := filepath.EvalSymlinks(parent)
	if err != nil {
		return fmt.Errorf("typed result path rejected")
	}
	relParent, err := filepath.Rel(resolvedRoot, resolvedParent)
	if err != nil || relParent == ".." || strings.HasPrefix(relParent, ".."+string(os.PathSeparator)) {
		return fmt.Errorf("typed result path rejected")
	}
	if info, err := os.Lstat(absDest); err == nil && info.Mode()&os.ModeSymlink != 0 {
		return fmt.Errorf("typed result path rejected")
	}
	return nil
}