iop/packages/go/agenttask/confinement_dispatch_test.go

package agenttask

import (
	"context"
	"errors"
	"reflect"
	"strings"
	"testing"
)

func TestValidatePreparedIsolationRequiresExactConfinementProof(t *testing.T) {
	projectID := ProjectID("project")
	workspaceID := WorkspaceID("workspace")
	project := ProjectRecord{
		ProjectID:   projectID,
		WorkspaceID: workspaceID,
		Intent: &StartIntent{
			ProjectID:        projectID,
			WorkspaceID:      workspaceID,
			ConfigRevision:   "config-r1",
			GrantRevision:    "grant-r1",
			WorkflowRevision: "workflow-r1",
		},
	}
	work := WorkRecord{
		Unit:      testUnit("work", WriteSetUnknown),
		AttemptID: "work#1",
	}
	target := ExecutionTarget{
		ProviderID:      "provider",
		ModelID:         "model",
		ProfileID:       "profile",
		ProfileRevision: "profile-r1",
		ConfigRevision:  "config-r1",
		Capacity:        1,
	}
	isolation := newFakeIsolation(t)
	prepared, err := isolation.Prepare(context.Background(), IsolationRequest{
		Project:        project,
		Work:           work,
		Target:         target,
		IdempotencyKey: "dispatch/project/work/1/isolation",
	})
	if err != nil {
		t.Fatalf("Prepare: %v", err)
	}
	if _, _, err := validatePreparedIsolation(project, work, target, prepared); err != nil {
		t.Fatalf("validate exact proof: %v", err)
	}

	missing := prepared
	missing.Confinement = nil
	if _, _, err := validatePreparedIsolation(project, work, target, missing); err == nil ||
		!strings.Contains(err.Error(), "incomplete strict ports") {
		t.Fatalf("missing proof error = %v", err)
	}

	tampered := prepared
	tamperedProof := *prepared.Confinement.(*fakeConfinement)
	tamperedProof.binding = prepared.Confinement.Binding()
	tamperedProof.binding.ProfileRevision = "profile-r2"
	tampered.Confinement = &tamperedProof
	if _, _, err := validatePreparedIsolation(project, work, target, tampered); err == nil ||
		!strings.Contains(err.Error(), "invalid executable confinement proof") {
		t.Fatalf("tampered proof error = %v", err)
	}
}

func TestManagerOwnsConfinementStartBeforeBindingProviderInvocation(t *testing.T) {
	commandType := reflect.TypeOf(ConfinementCommand{})
	for _, forbidden := range []string{"Stdin", "Stdout", "Stderr"} {
		if _, exists := commandType.FieldByName(forbidden); exists {
			t.Fatalf("ConfinementCommand still exposes caller-owned %s", forbidden)
		}
	}
	snapshot := testSnapshot("project", "workspace", testUnit("work", WriteSetUnknown))
	harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 1)
	harness.start("project", "workspace", nil)
	if err := harness.manager.Reconcile(context.Background()); err != nil {
		t.Fatalf("Reconcile: %v", err)
	}
	proof := harness.isolation.proof(dispatchKey("project", "work", 1) + "/isolation")
	if proof == nil {
		t.Fatal("isolation did not retain the executable proof")
	}
	if proof.startCount() != 1 {
		t.Fatalf("confinement starts = %d, want 1", proof.startCount())
	}
	commands := proof.startCommands()
	if len(commands) != 1 || commands[0].Name != "true" {
		t.Fatalf("confinement command = %#v, want exact prepared command", commands)
	}
	prepared, bound, proofStartsAtBind := harness.invoker.launchStats()
	if prepared != 1 || bound != 1 || len(proofStartsAtBind) != 1 || proofStartsAtBind[0] != 1 {
		t.Fatalf(
			"launch ordering prepare=%d bind=%d proof-starts-at-bind=%v, want 1/1/[1]",
			prepared,
			bound,
			proofStartsAtBind,
		)
	}
	proofHandles := proof.startedHandles()
	boundHandles := harness.invoker.startedHandles()
	if len(proofHandles) != 1 || len(boundHandles) != 1 ||
		proofHandles[0] != boundHandles[0] {
		t.Fatalf(
			"proof/bind handles = %#v/%#v, want one exact shared handle",
			proofHandles,
			boundHandles,
		)
	}
}

func TestManagerConfinementLaunchFailuresDoNotBindOrLeakChildren(t *testing.T) {
	tests := []struct {
		name        string
		configure   func(*managerHarness)
		wantStarts  int
		wantPrepare int
		wantBind    int
		wantChild   bool
		wantHandle  bool
		wantAbort   int
	}{
		{
			name: "prepare failure",
			configure: func(harness *managerHarness) {
				harness.invoker.prepareErr = errors.New("prepare failed")
			},
			wantPrepare: 1,
		},
		{
			name: "proof start failure",
			configure: func(harness *managerHarness) {
				harness.isolation.startErr = errors.New("proof start failed")
			},
			wantStarts:  1,
			wantPrepare: 1,
		},
		{
			name: "nil started handle",
			configure: func(harness *managerHarness) {
				harness.isolation.nilStarted = true
			},
			wantStarts:  1,
			wantPrepare: 1,
		},
		{
			name: "invalid started handle aborts child and pipes",
			configure: func(harness *managerHarness) {
				harness.isolation.invalidStart = true
				harness.invoker.command = ConfinementCommand{Name: "sleep", Args: []string{"5"}}
			},
			wantStarts:  1,
			wantPrepare: 1,
			wantChild:   true,
			wantHandle:  true,
			wantAbort:   1,
		},
		{
			name: "bind failure reaps child",
			configure: func(harness *managerHarness) {
				harness.invoker.bindErr = errors.New("bind failed")
				harness.invoker.command = ConfinementCommand{Name: "sleep", Args: []string{"5"}}
			},
			wantStarts:  1,
			wantPrepare: 1,
			wantBind:    1,
			wantChild:   true,
			wantHandle:  true,
			wantAbort:   1,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			snapshot := testSnapshot("project", "workspace", testUnit("work", WriteSetUnknown))
			harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 1)
			test.configure(harness)
			harness.start("project", "workspace", nil)
			if err := harness.manager.Reconcile(context.Background()); err != nil {
				t.Fatalf("Reconcile: %v", err)
			}
			proof := harness.isolation.proof(dispatchKey("project", "work", 1) + "/isolation")
			if proof == nil {
				t.Fatal("isolation did not retain the executable proof")
			}
			prepared, bound, _ := harness.invoker.launchStats()
			if proof.startCount() != test.wantStarts || prepared != test.wantPrepare || bound != test.wantBind {
				t.Fatalf(
					"starts/prepare/bind = %d/%d/%d, want %d/%d/%d",
					proof.startCount(),
					prepared,
					bound,
					test.wantStarts,
					test.wantPrepare,
					test.wantBind,
				)
			}
			children := proof.startedChildren()
			if test.wantChild {
				if len(children) != 1 || children[0].ProcessState == nil {
					t.Fatalf("bind failure leaked confined child: %#v", children)
				}
			} else if len(children) != 0 {
				t.Fatalf("failed launch unexpectedly started children: %#v", children)
			}
			handles := proof.startedHandles()
			if test.wantHandle {
				if len(handles) != 1 {
					t.Fatalf("started handles = %#v, want one", handles)
				}
				started, ok := handles[0].(*fakeStartedConfinement)
				if !ok {
					t.Fatalf("started handle type = %T", handles[0])
				}
				if started.abortCount() != test.wantAbort {
					t.Fatalf("handle aborts = %d, want %d", started.abortCount(), test.wantAbort)
				}
				if _, err := started.stdin.Write([]byte("leak")); err == nil {
					t.Fatal("partial-start cleanup left stdin open")
				}
				if _, err := started.stdout.Read(make([]byte, 1)); err == nil {
					t.Fatal("partial-start cleanup left stdout open")
				}
				if _, err := started.stderr.Read(make([]byte, 1)); err == nil {
					t.Fatal("partial-start cleanup left stderr open")
				}
			} else if len(handles) != 0 {
				t.Fatalf("failed launch unexpectedly returned handles: %#v", handles)
			}
			if harness.invoker.callCount() != 0 {
				t.Fatalf("failed launch submitted %d provider invocations", harness.invoker.callCount())
			}
		})
	}
}