Run every wrap into a macOS session inside that session

internal/api/app.go

@@ -124,6 +124,14 @@ type App struct {
 		// a notification, and the extension fetches asynchronously. Both
 		// wrap-init and exec block on this before letting a process start.
 		AwaitSnapshot(sessionID string, timeout time.Duration) bool
+		// AwaitRootDelivered blocks until the extension holds a snapshot of
+		// the session that names pid as a root. wrap-init needs it, because
+		// a session that the extension already holds passes AwaitSnapshot
+		// before the extension knows the new wrap process.
+		AwaitRootDelivered(sessionID string, pid int32, timeout time.Duration) bool
+		// RootPIDsForSession lists the session's roots. A non-empty list
+		// means that the session is already registered.
+		RootPIDsForSession(sessionID string) []int32
 	}
 
 	// acceptNotifyFDForTest, if non-nil, wraps the goroutine launch for
@@ -308,6 +316,8 @@ func (a *App) SetSessionTracker(t interface {
 	RegisterProcess(sessionID string, pid, ppid int32)
 	EndSession(sessionID string)
 	AwaitSnapshot(sessionID string, timeout time.Duration) bool
+	AwaitRootDelivered(sessionID string, pid int32, timeout time.Duration) bool
+	RootPIDsForSession(sessionID string) []int32
 }) {
 	a.sessionTracker = t
 }

internal/api/wrap_esf_darwin.go

@@ -54,6 +54,10 @@ func platformWrapInit(a *App, sessionID string, req types.WrapInitRequest) (type
 			fmt.Errorf("wrap: the agentmon system extension is not enforcing, so nothing would constrain this session (%s)", live.Detail)
 	}
 
+	// A second `agentmon wrap --session` joins a session that is already
+	// registered, and that the extension may already hold.
+	reused := len(a.sessionTracker.RootPIDsForSession(sessionID)) > 0
+
 	// ppid 0 marks this as a session root rather than a tracked child.
 	a.sessionTracker.RegisterProcess(sessionID, int32(req.CallerPID), 0)
 
@@ -71,6 +75,13 @@ func platformWrapInit(a *App, sessionID string, req types.WrapInitRequest) (type
 	// exec sessions then enforce nothing at all. That is a separate, known
 	// defect; this line is what wrap needs once it is fixed.
 	notifySessionRegistered()
+	// The session-registered notification makes the extension fetch only the
+	// sessions that it does not hold. For a reused session it must fetch
+	// again, to learn the new root. The policy-update notification makes it
+	// fetch every session it holds.
+	if reused {
+		notifyPolicyUpdated()
+	}
 
 	// Wait for the extension to actually hold the policy before letting the
 	// agent start.
@@ -82,12 +93,22 @@ func platformWrapInit(a *App, sessionID string, req types.WrapInitRequest) (type
 	// unenforced -- a `whoami` the policy denies succeeded, and so did a read of
 	// a denied file, while a curl a few hundred milliseconds later was blocked
 	// correctly. An agent's first actions are exactly the ones worth policing.
-	if !a.sessionTracker.AwaitSnapshot(sessionID, interceptionReadyTimeout) {
+	//
+	// The wait is for a snapshot that names this caller as a root. Any
+	// snapshot of the session is not enough: before, a second wrap into a
+	// session that the extension held returned at once, and its agent ran
+	// outside the session with no policy (SECENG-312).
+	if !a.sessionTracker.AwaitRootDelivered(sessionID, int32(req.CallerPID), interceptionReadyTimeout) {
 		// Undo the registration. A refused wrap-init must leave nothing behind:
 		// the server is telling the caller this session is not wrapped, and a
 		// session root left registered would have the extension attributing
 		// processes to a session the server considers unwrapped.
-		a.sessionTracker.EndSession(sessionID)
+		//
+		// A reused session stays. Another agent runs in it, and ending the
+		// session would stop the enforcement of that agent.
+		if !reused {
+			a.sessionTracker.EndSession(sessionID)
+		}
 		return types.WrapInitResponse{}, http.StatusServiceUnavailable,
 			fmt.Errorf("wrap: the system extension did not pick up this session's policy within %s, so the agent would start unconstrained", interceptionReadyTimeout)
 	}

internal/api/wrap_esf_darwin_test.go

@@ -25,6 +25,9 @@ type recordingTracker struct {
 	// picked up the policy. Zero value is "delivered", so tests about other
 	// things do not sit through the wait.
 	notDelivered bool
+
+	// awaited records the PIDs that AwaitRootDelivered waited for.
+	awaited []int32
 }
 
 func (r *recordingTracker) RegisterProcess(sessionID string, pid, ppid int32) {
@@ -42,6 +45,25 @@ func (r *recordingTracker) AwaitSnapshot(string, time.Duration) bool {
 	return !r.notDelivered
 }
 
+func (r *recordingTracker) AwaitRootDelivered(_ string, pid int32, _ time.Duration) bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.awaited = append(r.awaited, pid)
+	return !r.notDelivered
+}
+
+func (r *recordingTracker) RootPIDsForSession(sessionID string) []int32 {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	var out []int32
+	for _, reg := range r.registered {
+		if reg.sessionID == sessionID && reg.ppid == 0 {
+			out = append(out, reg.pid)
+		}
+	}
+	return out
+}
+
 func (r *recordingTracker) EndSession(sessionID string) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
@@ -226,3 +248,69 @@ func TestPlatformWrapInit_RefusesWhenSnapshotNeverArrives(t *testing.T) {
 		t.Errorf("refused wrap-init did not unregister the session; ended = %v", tr.ended)
 	}
 }
+
+// TestPlatformWrapInit_ReusedSessionWaitsForItsOwnRoot covers SECENG-312. A
+// second `agentmon wrap --session` joined a session that the extension
+// already held, so the wait for "a snapshot of the session" ended at once.
+// That snapshot did not name the new wrap process, and its agent ran outside
+// the session. wrap-init must wait for a snapshot that names the caller.
+func TestPlatformWrapInit_ReusedSessionWaitsForItsOwnRoot(t *testing.T) {
+	live := darwin.CheckSysExtLiveness()
+	if !live.Running {
+		t.Skip("no running system extension on this machine")
+	}
+	darwin.SetExtensionPresenceProbe(func() (bool, time.Time) { return true, time.Now() })
+	t.Cleanup(func() { darwin.SetExtensionPresenceProbe(nil) })
+
+	tr := &recordingTracker{}
+	a := &App{}
+	a.SetSessionTracker(tr)
+
+	for _, pid := range []int{4242, 5353} {
+		if _, _, err := platformWrapInit(a, "sess-1", types.WrapInitRequest{CallerPID: pid}); err != nil {
+			t.Fatalf("wrap-init for pid %d: %v", pid, err)
+		}
+	}
+
+	tr.mu.Lock()
+	defer tr.mu.Unlock()
+	if len(tr.awaited) != 2 || tr.awaited[0] != 4242 || tr.awaited[1] != 5353 {
+		t.Errorf("wrap-init waited for roots %v, want [4242 5353]", tr.awaited)
+	}
+}
+
+// TestPlatformWrapInit_ReusedSessionTimeoutKeepsSession: when the second
+// wrap into a session times out, the session must stay. Another agent runs
+// in it, and ending the session would stop the enforcement of that agent.
+func TestPlatformWrapInit_ReusedSessionTimeoutKeepsSession(t *testing.T) {
+	live := darwin.CheckSysExtLiveness()
+	if !live.Running {
+		t.Skip("no running system extension on this machine")
+	}
+	darwin.SetExtensionPresenceProbe(func() (bool, time.Time) { return true, time.Now() })
+	t.Cleanup(func() { darwin.SetExtensionPresenceProbe(nil) })
+
+	tr := &recordingTracker{}
+	a := &App{}
+	a.SetSessionTracker(tr)
+	if _, _, err := platformWrapInit(a, "sess-1", types.WrapInitRequest{CallerPID: 4242}); err != nil {
+		t.Fatalf("first wrap-init: %v", err)
+	}
+
+	tr.mu.Lock()
+	tr.notDelivered = true
+	tr.mu.Unlock()
+	_, code, err := platformWrapInit(a, "sess-1", types.WrapInitRequest{CallerPID: 5353})
+	if err == nil {
+		t.Fatal("second wrap-init succeeded although the extension never named its root")
+	}
+	if code != http.StatusServiceUnavailable {
+		t.Errorf("status = %d, want %d", code, http.StatusServiceUnavailable)
+	}
+
+	tr.mu.Lock()
+	defer tr.mu.Unlock()
+	if len(tr.ended) != 0 {
+		t.Errorf("a refused second wrap-init ended the session; ended = %v", tr.ended)
+	}
+}

internal/platform/darwin/policysock/handler.go

@@ -17,8 +17,9 @@ type SessionResolver interface {
 	SessionForPID(pid int32) string
 	LatestSession() (sessionID string, rootPID int32)
 	RootPIDForSession(sessionID string) int32
+	RootPIDsForSession(sessionID string) []int32
 	ActiveSessions() []string
-	NoteSnapshotDelivered(sessionID string)
+	NoteSnapshotDelivered(sessionID string, roots []int32)
 }
 
 // PolicyAdapter adapts the policy.Engine to the PolicyHandler interface.
@@ -330,16 +331,18 @@ func (a *PolicyAdapter) BuildPolicySnapshot(sessionID string, clientVersion uint
 	// two registrations; it fetches the latest and would never learn about the
 	// other. Handing back the full list lets it fetch what it is missing.
 	var activeSessions []string
+	var rootPIDs []int32
 	if a.sessions != nil {
 		activeSessions = a.sessions.ActiveSessions()
+		rootPIDs = a.sessions.RootPIDsForSession(sessionID)
 	}
 
 	// The extension asked for this session's policy and is about to receive it,
 	// so from here on it can enforce. wrap-init blocks on exactly this before
 	// letting the agent start; without it the agent's first commands run before
 	// the extension has any policy at all.
 	if a.sessions != nil {
-		a.sessions.NoteSnapshotDelivered(sessionID)
+		a.sessions.NoteSnapshotDelivered(sessionID, rootPIDs)
 	}
 
 	enforcement := networkEnforcement(p)
@@ -355,6 +358,7 @@ func (a *PolicyAdapter) BuildPolicySnapshot(sessionID string, clientVersion uint
 		Allow:              true,
 		SessionID:          sessionID,
 		RootPID:            rootPID,
+		RootPIDs:           rootPIDs,
 		SnapshotVersion:    nextSnapshotVersion(),
 		FileRules:          fileRules,
 		NetworkRules:       networkRules,

internal/platform/darwin/policysock/handler_test.go

@@ -387,7 +387,14 @@ func (m *mockSessionResolver) RootPIDForSession(sessionID string) int32 {
 	return 0
 }
 
-func (m *mockSessionResolver) NoteSnapshotDelivered(string) {}
+func (m *mockSessionResolver) RootPIDsForSession(sessionID string) []int32 {
+	if pid := m.RootPIDForSession(sessionID); pid != 0 {
+		return []int32{pid}
+	}
+	return nil
+}
+
+func (m *mockSessionResolver) NoteSnapshotDelivered(string, []int32) {}
 
 func (m *mockSessionResolver) ActiveSessions() []string {
 	seen := make(map[string]struct{}, len(m.sessions))

internal/platform/darwin/policysock/network_enforcement_test.go

@@ -222,7 +222,7 @@ func TestAwaitSnapshot(t *testing.T) {
 	t.Run("returns immediately once delivered", func(t *testing.T) {
 		tracker := NewSessionTracker()
 		tracker.RegisterProcess("session-a", 100, 0)
-		tracker.NoteSnapshotDelivered("session-a")
+		tracker.NoteSnapshotDelivered("session-a", nil)
 
 		start := time.Now()
 		if !tracker.AwaitSnapshot("session-a", time.Minute) {
@@ -240,7 +240,7 @@ func TestAwaitSnapshot(t *testing.T) {
 
 		go func() {
 			time.Sleep(50 * time.Millisecond)
-			tracker.NoteSnapshotDelivered("session-a")
+			tracker.NoteSnapshotDelivered("session-a", nil)
 		}()
 
 		if !tracker.AwaitSnapshot("session-a", 5*time.Second) {
@@ -279,7 +279,7 @@ func TestAwaitSnapshot(t *testing.T) {
 		}
 		// Giving up must not latch the session off: the extension may still
 		// arrive, and later commands should see it.
-		tracker.NoteSnapshotDelivered("session-a")
+		tracker.NoteSnapshotDelivered("session-a", nil)
 		if !tracker.AwaitSnapshot("session-a", time.Second) {
 			t.Error("a snapshot delivered after a timeout was ignored")
 		}
@@ -288,7 +288,7 @@ func TestAwaitSnapshot(t *testing.T) {
 	t.Run("ending a session clears both flags", func(t *testing.T) {
 		tracker := NewSessionTracker()
 		tracker.RegisterProcess("session-a", 100, 0)
-		tracker.NoteSnapshotDelivered("session-a")
+		tracker.NoteSnapshotDelivered("session-a", nil)
 		tracker.EndSession("session-a")
 
 		// A reused ID must not inherit the old session's readiness, or its
@@ -340,3 +340,40 @@ func TestNoSessionSnapshotShape(t *testing.T) {
 		t.Errorf("rule = %v, want no-policy-fail-closed", decoded["rule"])
 	}
 }
+
+// TestAwaitRootDelivered covers SECENG-312: a second wrap into a session
+// must wait for a snapshot that names its own process, not for any snapshot
+// of the session.
+func TestAwaitRootDelivered(t *testing.T) {
+	tracker := NewSessionTracker()
+	tracker.RegisterProcess("session-a", 100, 0)
+	tracker.NoteSnapshotDelivered("session-a", tracker.RootPIDsForSession("session-a"))
+
+	tracker.RegisterProcess("session-a", 200, 0)
+	tracker.RegisterProcess("session-a", 300, 100) // a child, not a root
+	if got := tracker.RootPIDsForSession("session-a"); len(got) != 2 || got[0] != 100 || got[1] != 200 {
+		t.Fatalf("RootPIDsForSession = %v, want [100 200]", got)
+	}
+	if !tracker.AwaitSnapshot("session-a", time.Millisecond) {
+		t.Fatal("AwaitSnapshot: the session has a delivered snapshot")
+	}
+	if tracker.AwaitRootDelivered("session-a", 200, 20*time.Millisecond) {
+		t.Fatal("AwaitRootDelivered reported pid 200 delivered; no snapshot named it")
+	}
+
+	go func() {
+		time.Sleep(20 * time.Millisecond)
+		tracker.NoteSnapshotDelivered("session-a", tracker.RootPIDsForSession("session-a"))
+	}()
+	if !tracker.AwaitRootDelivered("session-a", 200, 5*time.Second) {
+		t.Fatal("AwaitRootDelivered gave up on a snapshot that named pid 200")
+	}
+
+	tracker.EndSession("session-a")
+	if got := tracker.RootPIDsForSession("session-a"); got != nil {
+		t.Errorf("RootPIDsForSession after EndSession = %v, want nil", got)
+	}
+	if tracker.AwaitRootDelivered("session-a", 200, time.Millisecond) {
+		t.Error("AwaitRootDelivered reported a root of an ended session")
+	}
+}

internal/platform/darwin/policysock/protocol.go

@@ -140,9 +140,14 @@ type PolicyResponse struct {
 	// PendingApprovals answers list_approvals.
 	PendingApprovals []PendingApproval `json:"pending_approvals,omitempty"`
 
-	// Policy snapshot fields (returned by fetch_policy_snapshot)
+	// Policy snapshot fields (returned by fetch_policy_snapshot).
+	//
+	// RootPIDs lists every root of the session: root_pid first, then the
+	// process of each later `agentmon wrap --session` into it. The extension
+	// maps each one to the session.
 	SnapshotVersion uint64                `json:"version,omitempty"`
 	RootPID         int32                 `json:"root_pid,omitempty"`
+	RootPIDs        []int32               `json:"root_pids,omitempty"`
 	FileRules       []SnapshotFileRule    `json:"file_rules,omitempty"`
 	NetworkRules    []SnapshotNetworkRule `json:"network_rules,omitempty"`
 	DNSRules        []SnapshotDNSRule     `json:"dns_rules,omitempty"`