Implement UID-based cleanup for one-shot mode in procedure isolation

Add comprehensive cleanup functionality for setUID-isolated procedure runners:
- Clean up /tmp files owned by isolated UIDs after predictions complete
- Use os.Root for secure file operations to prevent path traversal
- Implement one-shot mode with graceful Stop() and ForceKill() fallback
- Add timeout-based cleanup with configurable CleanupTimeout
- Skip cleanup of workingdir/tmpdir when they're under /tmp
- Add test procedures and comprehensive test script
- Add CI job for cleanup testing alongside existing setuid tests

The cleanup ensures that files created by isolated procedure runners
are properly removed, preventing accumulation of orphaned files in
containerized environments like Docker and Kubernetes.

Author: tempusfrangit

.github/workflows/ci.yml

@@ -127,6 +127,22 @@ jobs:
       - run: ./script/init.sh
       - run: ./script/test-setuid.sh
 
+  test-setuid-cleanup:
+    name: Test Set UID cleanup in one-shot mode
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: astral-sh/setup-uv@v6
+      - run: ./script/init.sh
+      - name: Run setuid cleanup test
+        run: ./script/test-setuid-cleanup.sh
+        env:
+          PORT: 8080
+
   build-python:
     name: Build & verify python package
     runs-on: ubuntu-latest
@@ -178,6 +194,8 @@ jobs:
       - lint-go
       - test-go
       - test-python
+      - test-set-uid
+      - test-setuid-cleanup
     if: startsWith(github.ref, 'refs/tags/')
     runs-on: ubuntu-latest
     permissions:

cmd/cog/main.go

@@ -98,6 +98,7 @@ func buildServiceConfig(s *ServerCmd) (config.Config, error) {
 		MaxRunners:                s.MaxRunners,
 		RunnerShutdownGracePeriod: s.RunnerShutdownGracePeriod,
 		CleanupTimeout:            s.CleanupTimeout,
+		CleanupDirectories:        []string{"/tmp"},
 	}
 
 	log.Infow("service configuration",

internal/config/config.go

@@ -27,7 +27,8 @@ type Config struct {
 	RunnerShutdownGracePeriod time.Duration
 
 	// Cleanup configuration
-	CleanupTimeout time.Duration
+	CleanupTimeout     time.Duration
+	CleanupDirectories []string // Directories to walk for cleanup of files owned by isolated UIDs
 
 	// Environment configuration
 	EnvSet   map[string]string

internal/runner/manager.go

@@ -10,7 +10,6 @@ import (
 	"io/fs"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"runtime"
 	"sync"
 	"syscall"
@@ -348,12 +347,7 @@ func (m *Manager) createDefaultRunner(ctx context.Context) (*Runner, error) {
 		tmpDir:     tmpDir,
 		uploader:   uploader,
 	}
-	// Only enable forced shutdown for procedure mode
-	var forceShutdown *config.ForceShutdownSignal
-	if m.cfg.UseProcedureMode {
-		forceShutdown = m.cfg.ForceShutdown
-	}
-	runner, err := NewRunner(runtimeContext, runtimeCancel, runnerCtx, cmd, cogYaml.Concurrency.Max, m.cfg.CleanupTimeout, forceShutdown, m.baseLogger)
+	runner, err := NewRunner(runtimeContext, runtimeCancel, runnerCtx, cmd, cogYaml.Concurrency.Max, m.cfg, m.baseLogger)
 	if err != nil {
 		return nil, err
 	}
@@ -419,6 +413,35 @@ func (m *Manager) allocatePrediction(runner *Runner, req PredictionRequest) { //
 			delete(runner.pending, req.ID)
 			runner.mu.Unlock()
 
+			// In one-shot mode, stop runner after prediction completes to trigger cleanup
+			if m.cfg.OneShot && finalResponse.Status.IsCompleted() {
+				go func() {
+					logger := m.logger.Sugar()
+					logger.Infow("one-shot mode: stopping runner after prediction completion", "prediction_id", req.ID, "runner_id", runner.runnerCtx.id)
+
+					// Try graceful stop with timeout
+					stopDone := make(chan error, 1)
+					go func() {
+						stopDone <- runner.Stop()
+					}()
+
+					timeout := m.cfg.CleanupTimeout
+					if timeout == 0 {
+						timeout = 10 * time.Second // Default timeout
+					}
+
+					select {
+					case err := <-stopDone:
+						if err != nil {
+							logger.Errorw("failed to stop runner in one-shot mode", "error", err, "runner_id", runner.runnerCtx.id)
+						}
+					case <-time.After(timeout):
+						logger.Warnw("stop timeout exceeded in one-shot mode, falling back to force kill", "timeout", timeout, "runner_id", runner.runnerCtx.id)
+						runner.ForceKill()
+					}
+				}()
+			}
+
 			if cancel != nil {
 				cancel()
 			}
@@ -626,20 +649,28 @@ func (m *Manager) createProcedureRunner(runnerName, procedureHash string) (*Runn
 	env = append(env, "TMPDIR="+tmpDir)
 	cmd.Env = env
 
-	// Apply setUID isolation for procedure runners if needed
+	var allocatedUID *int
 	if m.shouldUseSetUID() {
 		uid, err := AllocateUID()
 		if err != nil {
 			runtimeCancel()
 			return nil, fmt.Errorf("failed to allocate UID: %w", err)
 		}
+		allocatedUID = &uid
+
+		// Use os.Root for secure ownership changes
+		workingRoot, err := os.OpenRoot(workingDir)
+		if err != nil {
+			runtimeCancel()
+			return nil, fmt.Errorf("failed to open working directory root: %w", err)
+		}
+		defer func() { _ = workingRoot.Close() }()
 
-		// Change ownership of source directory (workingDir)
-		err = filepath.WalkDir(workingDir, func(path string, d fs.DirEntry, err error) error {
+		err = fs.WalkDir(workingRoot.FS(), ".", func(path string, d fs.DirEntry, err error) error {
 			if err != nil {
 				return err
 			}
-			if lchownErr := os.Lchown(path, uid, NoGroupGID); lchownErr != nil {
+			if lchownErr := workingRoot.Lchown(path, uid, NoGroupGID); lchownErr != nil {
 				log.Errorw("failed to change ownership", "path", path, "uid", uid, "error", lchownErr)
 				return lchownErr
 			}
@@ -650,19 +681,24 @@ func (m *Manager) createProcedureRunner(runnerName, procedureHash string) (*Runn
 			return nil, fmt.Errorf("failed to change ownership of source directory: %w", err)
 		}
 
-		// Make working dir writable by unprivileged Python process
-		if err := os.Lchown(workingDir, uid, NoGroupGID); err != nil {
+		if err := workingRoot.Lchown(".", uid, NoGroupGID); err != nil {
 			log.Errorw("failed to change ownership of working directory", "path", workingDir, "uid", uid, "error", err)
 			runtimeCancel()
 			return nil, fmt.Errorf("failed to change ownership of working directory: %w", err)
 		}
-		// Change ownership of temp directory
-		if err := os.Lchown(tmpDir, uid, NoGroupGID); err != nil {
+
+		tmpRoot, err := os.OpenRoot(tmpDir)
+		if err != nil {
+			runtimeCancel()
+			return nil, fmt.Errorf("failed to open temp directory root: %w", err)
+		}
+		defer func() { _ = tmpRoot.Close() }()
+
+		if err := tmpRoot.Lchown(".", uid, NoGroupGID); err != nil {
 			log.Errorw("failed to change ownership of temp directory", "path", tmpDir, "uid", uid, "error", err)
 			runtimeCancel()
 			return nil, fmt.Errorf("failed to change ownership of temp directory: %w", err)
 		}
-		// Use syscall.Credential to run process as unprivileged user from start
 		cmd.SysProcAttr.Credential = &syscall.Credential{
 			Uid: uint32(uid), //nolint:gosec // this is guarded in isolation .allocate, cannot exceed const MaxUID
 			Gid: uint32(NoGroupGID),
@@ -675,19 +711,17 @@ func (m *Manager) createProcedureRunner(runnerName, procedureHash string) (*Runn
 	if m.cfg.UploadURL != "" {
 		uploader = newUploader(m.cfg.UploadURL)
 	}
+
 	runnerCtx := RunnerContext{
-		id:         runnerName,
-		workingdir: workingDir,
-		tmpDir:     tmpDir,
-		uploader:   uploader,
+		id:                 runnerName,
+		workingdir:         workingDir,
+		tmpDir:             tmpDir,
+		uploader:           uploader,
+		uid:                allocatedUID,
+		cleanupDirectories: m.cfg.CleanupDirectories,
 	}
 
-	// Only enable forced shutdown for procedure mode
-	var forceShutdown *config.ForceShutdownSignal
-	if m.cfg.UseProcedureMode {
-		forceShutdown = m.cfg.ForceShutdown
-	}
-	runner, err := NewRunner(runtimeContext, runtimeCancel, runnerCtx, cmd, 1, m.cfg.CleanupTimeout, forceShutdown, m.baseLogger)
+	runner, err := NewRunner(runtimeContext, runtimeCancel, runnerCtx, cmd, 1, m.cfg, m.baseLogger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create runner: %w", err)
 	}

internal/runner/runner.go

@@ -983,7 +983,7 @@ func verifyProcessGroupTerminated(pid int) error {
 }
 
 // NewRunner creates a new runner instance with the given context
-func NewRunner(ctx context.Context, ctxCancel context.CancelFunc, runnerCtx RunnerContext, command *exec.Cmd, maxConcurrency int, cleanupTimeout time.Duration, forceShutdown *config.ForceShutdownSignal, logger *zap.Logger) (*Runner, error) {
+func NewRunner(ctx context.Context, ctxCancel context.CancelFunc, runnerCtx RunnerContext, command *exec.Cmd, maxConcurrency int, cfg config.Config, logger *zap.Logger) (*Runner, error) {
 	if maxConcurrency <= 0 {
 		maxConcurrency = 1
 	}
@@ -1005,8 +1005,8 @@ func NewRunner(ctx context.Context, ctxCancel context.CancelFunc, runnerCtx Runn
 		readyForShutdown:   make(chan struct{}),
 		setupComplete:      make(chan struct{}),
 		logCaptureComplete: make(chan struct{}),
-		cleanupTimeout:     cleanupTimeout,
-		forceShutdown:      forceShutdown,
+		cleanupTimeout:     cfg.CleanupTimeout,
+		forceShutdown:      cfg.ForceShutdown,
 		logger:             runnerLogger,
 	}
 

internal/runner/runner_test.go

@@ -789,7 +789,8 @@ func TestNewRunner(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
-		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, 0, nil, zaptest.NewLogger(t))
+		cfg := config.Config{}
+		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, cfg, zaptest.NewLogger(t))
 		require.NoError(t, err)
 
 		assert.Equal(t, "test-runner", r.runnerCtx.id)
@@ -825,7 +826,8 @@ func TestNewRunner(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
-		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, 0, nil, zaptest.NewLogger(t))
+		cfg := config.Config{}
+		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, cfg, zaptest.NewLogger(t))
 		require.NoError(t, err)
 
 		// Should store the command correctly
@@ -849,7 +851,8 @@ func TestNewRunner(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
-		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, 0, nil, zaptest.NewLogger(t))
+		cfg := config.Config{}
+		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, cfg, zaptest.NewLogger(t))
 		require.NoError(t, err)
 		require.NotNil(t, r)
 
@@ -887,7 +890,8 @@ func TestProcedureRunnerCreation(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
-		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, 0, nil, zaptest.NewLogger(t))
+		cfg := config.Config{}
+		r, err := NewRunner(ctx, cancel, runnerCtx, cmd, 1, cfg, zaptest.NewLogger(t))
 		require.NoError(t, err)
 
 		assert.Equal(t, "proc-runner", r.runnerCtx.id)

internal/runner/types.go

@@ -7,10 +7,13 @@ import (
 	"encoding/base32"
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"os"
+	"path/filepath"
 	"strings"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/replicate/cog-runtime/internal/util"
@@ -210,16 +213,98 @@ func (r RunnerID) String() string {
 
 // RunnerContext contains everything a runner needs to operate
 type RunnerContext struct {
-	id         string
-	workingdir string
-	tmpDir     string
-	uploader   *uploader
+	id                 string
+	workingdir         string
+	tmpDir             string
+	uploader           *uploader
+	uid                *int     // UID used for setUID isolation, nil if not using setUID
+	cleanupDirectories []string // Directories to walk for cleanup of files owned by isolated UIDs
 }
 
 func (rc *RunnerContext) Cleanup() error {
 	if rc.tmpDir != "" {
-		return os.RemoveAll(rc.tmpDir)
+		if err := os.RemoveAll(rc.tmpDir); err != nil {
+			return err
+		}
+	}
+
+	// Clean up files in configured directories owned by this UID when using setUID isolation
+	if rc.uid != nil && len(rc.cleanupDirectories) > 0 {
+		return rc.cleanupDirectoriesFiles()
+	}
+
+	return nil
+}
+
+// cleanupDirectoriesFiles removes files in configured directories owned by the isolated UID
+func (rc *RunnerContext) cleanupDirectoriesFiles() error {
+	if rc.uid == nil {
+		return nil
 	}
+
+	// Avoid cleaning our own workingdir/tmpdir if they're in the cleanup directories
+	skipPaths := make(map[string]bool)
+	for _, cleanupDir := range rc.cleanupDirectories {
+		if strings.HasPrefix(rc.workingdir, cleanupDir+"/") {
+			skipPaths[rc.workingdir] = true
+		}
+		if strings.HasPrefix(rc.tmpDir, cleanupDir+"/") {
+			skipPaths[rc.tmpDir] = true
+		}
+	}
+
+	for _, cleanupDir := range rc.cleanupDirectories {
+		// Use os.OpenRoot to create a secure chrooted view of the cleanup directory
+		root, err := os.OpenRoot(cleanupDir)
+		if err != nil {
+			continue // Skip directories we can't root into
+		}
+
+		err = fs.WalkDir(root.FS(), ".", func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return nil // Continue walking on errors
+			}
+
+			// Convert relative path back to absolute for skipPaths check
+			absPath := filepath.Join(cleanupDir, path)
+			if skipPaths[absPath] {
+				return filepath.SkipDir
+			}
+
+			// Don't follow symlinks
+			if d.Type()&fs.ModeSymlink != 0 {
+				return nil
+			}
+
+			// Check if file is owned by our UID using root.Stat
+			info, err := root.Stat(path)
+			if err != nil {
+				return nil // Continue on stat errors
+			}
+
+			if stat, ok := info.Sys().(*syscall.Stat_t); ok {
+				if int(stat.Uid) == *rc.uid {
+					if err := root.RemoveAll(path); err != nil {
+						// Log error but continue cleanup
+						return nil
+					}
+					if d.IsDir() {
+						return filepath.SkipDir
+					}
+				}
+			}
+
+			return nil
+		})
+
+		// Close the root after processing this directory
+		_ = root.Close()
+
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }