iop/apps/node/internal/workspace/cleanup.go

package workspace

import (
	"context"
	"errors"
	"path"
	"sort"
	"strings"
	"time"

	iop "iop/proto/gen/iop"
)

var errCleanupUnsupported = errors.New("workspace cleanup is unsupported on this platform")

// WriteInternalArtifact creates a new Node-owned request artifact. The caller
// supplies only a path relative to its immutable request namespace; the public
// workspace tool surface cannot invoke this helper or name .iop directly.
func (r *Runtime) WriteInternalArtifact(requestID, relativePath string, content []byte) error {
	if len(content) > maxInternalArtifactSize {
		return ErrInvalidRequest
	}
	req, err := r.Request(requestID)
	if err != nil {
		return err
	}
	name, err := internalArtifactPath(relativePath)
	if err != nil {
		return ErrInvalidRequest
	}
	req.mu.Lock()
	defer req.mu.Unlock()
	if req.cleaning {
		return ErrClosed
	}
	created, err := createOwnedArtifact(req.entry, req.internalPrefix, name, content, req.artifacts)
	if err != nil {
		return err
	}
	if len(req.artifacts)+len(created) > maxCleanupArtifacts {
		rollbackOwnedArtifacts(req.entry, created)
		return ErrInvalidRequest
	}
	for _, artifact := range created {
		req.artifacts[artifact.relative] = artifact
	}
	return nil
}

func internalArtifactPath(value string) (string, error) {
	if value == "" || len(value) > 1024 || value == "." || path.IsAbs(value) || path.Clean(value) != value || strings.Contains(value, "\\") || strings.ContainsRune(value, 0) || strings.HasPrefix(value, "../") || value == ".." {
		return "", errInvalidPath
	}
	return value, nil
}

// Cleanup elects one result owner for a request, cancels all of its command
// groups, validates the exact request tree against the in-memory ownership
// inventory, and removes only matching entries with non-recursive operations.
func (r *Runtime) Cleanup(ctx context.Context, requestID string) CleanupResult {
	if !validRequestID(requestID) {
		return cleanupFailure(iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST)
	}
	if ctx == nil {
		ctx = context.Background()
	}

	r.cleanupMu.Lock()
	if existing := r.cleanupCalls[requestID]; existing != nil {
		r.cleanupMu.Unlock()
		<-existing.done
		return existing.result
	}
	r.mu.RLock()
	req := r.requests[requestID]
	r.mu.RUnlock()
	if req == nil {
		r.cleanupMu.Unlock()
		return cleanupFailure(iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_FOUND)
	}
	call := &cleanupCall{done: make(chan struct{})}
	r.cleanupCalls[requestID] = call
	r.cleanupMu.Unlock()

	startedAt := time.Now()
	call.result = r.performCleanup(ctx, req)
	r.observeCleanup(req, call.result, time.Since(startedAt).Milliseconds())
	close(call.done)
	r.cleanupMu.Lock()
	r.cleanupOrder = append(r.cleanupOrder, requestID)
	for len(r.cleanupOrder) > completedCleanupLimit {
		evicted := r.cleanupOrder[0]
		r.cleanupOrder = r.cleanupOrder[1:]
		delete(r.cleanupCalls, evicted)
	}
	r.cleanupMu.Unlock()
	return call.result
}

func (r *Runtime) performCleanup(ctx context.Context, req *Request) CleanupResult {
	req.mu.Lock()
	req.cleaning = true
	artifacts := make(map[string]ownedArtifact, len(req.artifacts))
	for relative, artifact := range req.artifacts {
		artifacts[relative] = artifact
	}
	ownedParents := append([]ownedArtifact(nil), req.ownedParents...)
	req.mu.Unlock()

	executions := r.cancelRequestCommands(req.id)
	cleanupCtx, cancel := boundedCleanupContext(ctx)
	defer cancel()
	if !waitForCommandCleanup(cleanupCtx, executions) {
		r.closeRequestAuthority(req)
		return CleanupResult{
			Status:           iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT,
			Code:             iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT,
			CleanedProcesses: int32(len(executions)),
		}
	}
	cleanedProcesses := int32(len(executions))
	cleanedArtifacts, err := validateAndRemoveOwnedArtifacts(req.entry, req.internalPrefix, artifacts, ownedParents)
	r.closeRequestAuthority(req)
	if err != nil {
		if errors.Is(err, errCleanupUnsupported) {
			return CleanupResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSUPPORTED, CleanedProcesses: cleanedProcesses}
		}
		return CleanupResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, CleanedProcesses: cleanedProcesses}
	}
	return CleanupResult{
		Status:           iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
		Code:             iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSPECIFIED,
		CleanedProcesses: cleanedProcesses,
		CleanedArtifacts: int32(cleanedArtifacts),
	}
}

func (r *Runtime) cancelRequestCommands(requestID string) []*commandExecution {
	r.commandsMu.Lock()
	defer r.commandsMu.Unlock()
	executions := make([]*commandExecution, 0)
	for key, execution := range r.activeCommands {
		if key.requestID != requestID || execution == nil {
			continue
		}
		if execution.requestCancel() {
			executions = append(executions, execution)
		}
	}
	return executions
}

func boundedCleanupContext(parent context.Context) (context.Context, context.CancelFunc) {
	if deadline, ok := parent.Deadline(); ok && time.Until(deadline) <= defaultCleanupTimeout {
		return context.WithCancel(parent)
	}
	return context.WithTimeout(parent, defaultCleanupTimeout)
}

func waitForCommandCleanup(ctx context.Context, executions []*commandExecution) bool {
	for _, execution := range executions {
		select {
		case <-execution.done:
		case <-ctx.Done():
			return false
		}
	}
	return true
}

func (r *Runtime) closeRequestAuthority(req *Request) {
	r.mu.Lock()
	if r.requests[req.id] == req {
		delete(r.requests, req.id)
	}
	r.mu.Unlock()
	r.commandsMu.Lock()
	for key := range r.cancelledCommands {
		if key.requestID == req.id {
			delete(r.cancelledCommands, key)
		}
	}
	r.commandsMu.Unlock()
}

func cleanupFailure(status iop.WorkspaceStatus, code iop.WorkspaceErrorCode) CleanupResult {
	return CleanupResult{Status: status, Code: code}
}

func sortedArtifactsDeepestFirst(artifacts map[string]ownedArtifact) []ownedArtifact {
	ordered := make([]ownedArtifact, 0, len(artifacts))
	for _, artifact := range artifacts {
		ordered = append(ordered, artifact)
	}
	sort.Slice(ordered, func(i, j int) bool {
		leftDepth := strings.Count(ordered[i].relative, "/")
		rightDepth := strings.Count(ordered[j].relative, "/")
		if leftDepth != rightDepth {
			return leftDepth > rightDepth
		}
		return ordered[i].relative > ordered[j].relative
	})
	return ordered
}