Agent-System/internal/workspace/capability.go

package workspace

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"chat-agent-system/internal/app/uow"
)

// Capability is the only module allowed to turn a WorkspaceRef into bytes on disk.
type Capability struct {
	dataRoot string
	leases   *LeaseTable
}

func NewCapability(dataRoot string, leases *LeaseTable) *Capability {
	if leases == nil {
		leases = NewLeaseTable()
	}
	return &Capability{dataRoot: dataRoot, leases: leases}
}

func (c *Capability) Issue(ref WorkspaceRef) (WorkspaceRef, error) {
	if err := c.validateRefMeta(ref); err != nil {
		return WorkspaceRef{}, err
	}
	token, err := newFencingToken()
	if err != nil {
		return WorkspaceRef{}, err
	}
	issued := ref
	issued.SchemaVersion = SchemaVersion
	issued.FencingToken = token
	c.leases.Issue(issued)
	return issued, nil
}

func (c *Capability) ReadArtifact(ctx context.Context, ref WorkspaceRef, key string) ([]byte, error) {
	if err := c.guard(ctx, ref, false); err != nil {
		return nil, err
	}
	path, err := c.resolveKey(ref, key)
	if err != nil {
		return nil, err
	}
	resolved, err := c.enforceRealpath(path)
	if err != nil {
		return nil, err
	}
	return os.ReadFile(resolved)
}

func (c *Capability) WriteArtifact(ctx context.Context, ref WorkspaceRef, key string, body []byte) error {
	if err := c.guard(ctx, ref, true); err != nil {
		return err
	}
	path, err := c.resolveKey(ref, key)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	tmp := path + ".tmp-" + ref.FencingToken[:8]
	if err := os.WriteFile(tmp, body, 0o644); err != nil {
		return err
	}
	if _, err := c.enforceRealpath(tmp); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	return nil
}

func (c *Capability) CommitProjection(ctx context.Context, ref WorkspaceRef, key, digest string) error {
	if err := c.guard(ctx, ref, true); err != nil {
		return err
	}
	body, err := c.ReadArtifact(ctx, ref, key)
	if err != nil {
		return err
	}
	sum := sha256.Sum256(body)
	got := hex.EncodeToString(sum[:])
	if digest != "" && digest != got {
		return fmt.Errorf("workspace: digest mismatch")
	}
	marker := []byte(got + "\ncomplete\n")
	return c.WriteArtifact(ctx, ref, key+".complete", marker)
}

func (c *Capability) guard(ctx context.Context, ref WorkspaceRef, write bool) error {
	if ctx != nil {
		if _, ok := uow.ScopeFrom(ctx); ok {
			return ErrIOInTransaction
		}
	}
	if err := c.validateRefMeta(ref); err != nil {
		return err
	}
	if write {
		return c.leases.Assert(ref)
	}
	return nil
}

func (c *Capability) validateRefMeta(ref WorkspaceRef) error {
	if strings.TrimSpace(ref.ID) == "" {
		return ErrInvalidKey
	}
	if ref.SchemaVersion != 0 && ref.SchemaVersion != SchemaVersion {
		return ErrUnknownSchema
	}
	if ref.ProjectScope < 0 {
		return ErrCrossProject
	}
	return nil
}

func (c *Capability) resolveKey(ref WorkspaceRef, key string) (string, error) {
	if err := validateArtifactKey(key); err != nil {
		return "", err
	}
	root := filepath.Join(c.dataRoot, "workspaces", ref.ID)
	joined := filepath.Join(root, filepath.FromSlash(key))
	absRoot, err := filepath.Abs(root)
	if err != nil {
		return "", err
	}
	absJoin, err := filepath.Abs(joined)
	if err != nil {
		return "", err
	}
	rel, err := filepath.Rel(absRoot, absJoin)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return "", ErrPathEscape
	}
	return absJoin, nil
}

func (c *Capability) enforceRealpath(path string) (string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return path, nil
		}
		return "", err
	}
	if info.Mode()&os.ModeSymlink != 0 {
		return "", ErrPathEscape
	}
	if info.Mode()&os.ModeNamedPipe != 0 {
		return "", ErrPathEscape
	}
	resolved, err := filepath.EvalSymlinks(filepath.Dir(path))
	if err != nil {
		return "", err
	}
	full := filepath.Join(resolved, filepath.Base(path))
	root, err := filepath.Abs(c.dataRoot)
	if err != nil {
		return "", err
	}
	rel, err := filepath.Rel(root, full)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return "", ErrPathEscape
	}
	if st, ok := info.Sys().(*syscall.Stat_t); ok && st.Nlink > 1 && !info.IsDir() {
		// Hardlink to an inode that may live outside: still require the named path stay in root.
		if !strings.HasPrefix(full, root) {
			return "", ErrPathEscape
		}
	}
	return full, nil
}

func validateArtifactKey(key string) error {
	key = strings.TrimSpace(key)
	if key == "" {
		return ErrInvalidKey
	}
	if strings.ContainsRune(key, 0) {
		return ErrInvalidKey
	}
	if filepath.IsAbs(key) || strings.HasPrefix(key, "/") || strings.HasPrefix(key, `\`) {
		return ErrInvalidKey
	}
	lower := strings.ToLower(key)
	if strings.Contains(lower, "%2e%2e") || strings.Contains(lower, "%2f") || strings.Contains(lower, "%5c") {
		return ErrInvalidKey
	}
	cleaned := filepath.ToSlash(key)
	for _, part := range strings.Split(cleaned, "/") {
		if part == "" || part == "." || part == ".." {
			return ErrInvalidKey
		}
	}
	return nil
}

func newFencingToken() (string, error) {
	var b [16]byte
	if _, err := io.ReadFull(rand.Reader, b[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(b[:]), nil
}