fix: potential double free after use bug.

Author: Teryl Taylor

crates/cpex-ffi/src/lib.rs

@@ -706,7 +706,29 @@ pub unsafe extern "C" fn cpex_plugin_names(
 /// Returns MessagePack-encoded PipelineResult + opaque handles for
 /// context table and background tasks.
 ///
-/// Returns 0 on success, -1 on failure.
+/// # Ownership contract
+///
+/// **The caller's input `context_table` is unconditionally consumed
+/// by this function** — even on error paths (RC_INVALID_HANDLE,
+/// RC_INVALID_INPUT, RC_PARSE_ERROR, RC_TIMEOUT, RC_PANIC, etc.).
+/// The Box is freed inside `cpex_invoke`; the caller's pointer is
+/// dead once this function returns. This mirrors the pattern used
+/// by `cpex_wait_background` and lets the Go binding nil its handle
+/// unconditionally after the call without leaking the underlying Box.
+///
+/// On `RC_OK`, a **fresh** `CpexContextTableInner` Box is allocated
+/// and its raw pointer is written to `*context_table_out`. On any
+/// non-OK return, `*context_table_out` is left as a null pointer
+/// (initialized at function entry). The other out parameters
+/// (`result_msgpack_out`, `result_len_out`, `bg_handle_out`) follow
+/// the same discipline: null/zero on error, populated on success.
+///
+/// Pre-P0-1 the function consumed the input only after validation
+/// passed but before `run_safely`. On `RC_TIMEOUT` / `RC_PANIC` the
+/// input had been consumed but `*context_table_out` was never written,
+/// so the Go wrapper kept its stale handle and a subsequent
+/// `ContextTable.Close()` ran `cpex_release_context_table` on
+/// already-freed memory.
 ///
 /// # Safety
 /// All pointer parameters must be valid or NULL where documented.
@@ -731,7 +753,43 @@ pub unsafe extern "C" fn cpex_invoke(
     context_table_out: *mut *mut CpexContextTableInner,
     bg_handle_out: *mut *mut CpexBackgroundTasksInner,
 ) -> c_int {
-    // Validate manager handle
+    // Initialize all out params to safe defaults. Any early return
+    // from here on leaves a consistent state for the caller: every
+    // out pointer is null/zero, so a downstream attempt to dereference
+    // produces a clean null-deref crash rather than reading uninit
+    // stack memory. The success path overwrites these at the end.
+    *result_msgpack_out = std::ptr::null_mut();
+    *result_len_out = 0;
+    *context_table_out = std::ptr::null_mut();
+    *bg_handle_out = std::ptr::null_mut();
+
+    // Take ownership of the input context_table *immediately*, before
+    // any validation that could return an error code. From this point
+    // on, the caller's `context_table` pointer is dead — equivalent
+    // to free'd memory from the caller's perspective. This mirrors
+    // how `cpex_wait_background` handles `bg_handle`: ownership
+    // transfers on entry, the caller nils its reference, and Rust
+    // is responsible for the Box's lifetime from then on. Pre-fix,
+    // consumption happened mid-function after some validations, which
+    // meant validation errors left the input alive (one ownership
+    // model) and post-validation errors left it consumed without
+    // writing `*context_table_out` (a *different* ownership model).
+    // Two contracts in one function is exactly what produced the
+    // P0-1 UAF.
+    let input_ctx_table: Option<PluginContextTable> = if context_table.is_null() {
+        None
+    } else {
+        // Box::from_raw consumes the allocation; it'll drop at the
+        // end of this scope if not moved into Some(...). When moved
+        // into Some(...), the table value lives until invoke_by_name
+        // either uses it or it's dropped on a Future-cancellation
+        // path (RC_TIMEOUT). Either way the Box is gone.
+        let ct = Box::from_raw(context_table);
+        Some(ct.table)
+    };
+
+    // Validate manager handle. `input_ctx_table` already owns the
+    // input data — if we return here, it drops cleanly.
     let inner = match mgr.as_ref() {
         Some(m) => m,
         None => return RC_INVALID_HANDLE,
@@ -774,22 +832,17 @@ pub unsafe extern "C" fn cpex_invoke(
         Extensions::default()
     };
 
-    // Get or create context table
-    let ctx_table: Option<PluginContextTable> = if context_table.is_null() {
-        None
-    } else {
-        let ct = Box::from_raw(context_table);
-        Some(ct.table)
-    };
-
     // Invoke the hook with wall-clock timeout + panic catch.
     let (mut result, bg) = match run_safely(
         inner
             .manager
-            .invoke_by_name(name, payload, extensions, ctx_table),
+            .invoke_by_name(name, payload, extensions, input_ctx_table),
         "cpex_invoke",
     ) {
         SafeRun::Ok(r) => r,
+        // *context_table_out is already null (set at function entry);
+        // the input table has been consumed by invoke_by_name's call
+        // frame and dropped. Caller's handle is dead, no replacement.
         other => return other.rc(), // RC_TIMEOUT or RC_PANIC; already logged
     };
 

go/cpex/manager.go

@@ -341,19 +341,18 @@ func (m *PluginManager) InvokeByName(
 	cHookName := C.CString(hookName)
 	defer C.free(unsafe.Pointer(cHookName))
 
-	// Pass the context-table handle to Rust but DO NOT nil our local
-	// reference until we know Rust succeeded. Rust consumes the handle
-	// only at the moment of invoke (after all input validation), so
-	// pre-invoke failures (bad payload, bad extensions, etc.) leave
-	// the handle untouched and the caller's ContextTable remains valid.
-	//
-	// Caveat: on a post-invoke failure (rare — only result-serialization
-	// OOM), Rust has consumed the box but doesn't write ctOut, so the
-	// caller's ContextTable handle becomes dangling. The caller should
-	// not reuse a ContextTable after an InvokeByName error.
+	// Pass the context-table handle to Rust. Per the post-P0-1 FFI
+	// contract, `cpex_invoke` takes ownership of `ctHandle`
+	// UNCONDITIONALLY on entry — same pattern as `cpex_wait_background`
+	// with `bg_handle`. We nil our local reference immediately so the
+	// caller's `ContextTable` can't be accidentally reused after this
+	// call (its underlying Box is gone regardless of the eventual rc).
+	// On RC_OK a fresh handle lands in `ctOut`; on any error path
+	// `ctOut` stays nil and the caller's context-table chain ends here.
 	var ctHandle C.CpexContextTable
 	if contextTable != nil {
 		ctHandle = contextTable.handle
+		contextTable.handle = nil
 	}
 
 	var resultPtr *C.uint8_t
@@ -386,16 +385,13 @@ func (m *PluginManager) InvokeByName(
 	)
 
 	if rc != 0 {
+		// `ctOut` is null on every non-OK return per the post-P0-1
+		// contract. The caller's `contextTable.handle` is already nil
+		// (we cleared it above before the call), so there's no
+		// dangling-handle risk on error paths.
 		return nil, nil, nil, errorFromRC(int(rc), "InvokeByName")
 	}
 
-	// Rust succeeded — it consumed ctHandle and produced ctOut.
-	// NOW it's safe to nil the caller's reference (the original Box
-	// was consumed by Rust; its successor is in ctOut).
-	if contextTable != nil {
-		contextTable.handle = nil
-	}
-
 	// Deserialize result from MessagePack
 	resultBytes := C.GoBytes(unsafe.Pointer(resultPtr), resultLen)
 	C.cpex_free_bytes((*C.uint8_t)(unsafe.Pointer(resultPtr)), resultLen)

go/cpex/manager_test.go

@@ -1173,3 +1173,113 @@ func TestLoadConfigInvalidYAML(t *testing.T) {
 		t.Error("expected error for invalid YAML")
 	}
 }
+
+// TestInvokeByNameErrorDoesNotUAFContextTable is the regression guard
+// for P0-1. Pre-fix, `cpex_invoke` consumed the input ContextTable's
+// Box mid-function but didn't write *context_table_out on
+// RC_TIMEOUT / RC_PANIC / RC_PARSE_ERROR — the Go wrapper kept its
+// stale handle and a subsequent Close() called
+// cpex_release_context_table on already-freed memory.
+//
+// Post-fix, the Go wrapper nils its `contextTable.handle` immediately
+// after handing it to Rust (mirroring the bg_handle pattern), so even
+// if Rust errors out without producing a replacement, no dangling
+// handle survives.
+//
+// This test:
+//   1. Performs a successful invoke to get a real ContextTable.
+//   2. Calls InvokeByName again with that ContextTable PLUS an
+//      invalid payload_type that forces Rust to return RC_PARSE_ERROR
+//      AFTER the consumption point.
+//   3. Confirms the second call errored (sanity).
+//   4. Calls Close() on the original ContextTable — must NOT crash.
+//      Pre-fix this was a UAF (free of already-freed memory).
+func TestInvokeByNameErrorDoesNotUAFContextTable(t *testing.T) {
+	mgr, err := NewPluginManagerDefault()
+	if err != nil {
+		t.Fatalf("NewPluginManagerDefault failed: %v", err)
+	}
+	defer mgr.Shutdown()
+	if err := mgr.Initialize(); err != nil {
+		t.Fatalf("Initialize failed: %v", err)
+	}
+
+	// 1. Successful first invoke — gives us a real, Rust-allocated
+	//    ContextTable. Calling Close() on this pre-fix would have
+	//    been the second free.
+	payload := map[string]any{"tool_name": "test"}
+	_, ctxTable, bg, err := mgr.InvokeByName("hook1", PayloadGeneric, payload, &Extensions{}, nil)
+	if err != nil {
+		t.Fatalf("first invoke failed: %v", err)
+	}
+	bg.Close()
+	if ctxTable == nil || ctxTable.handle == nil {
+		t.Fatal("expected a non-nil ContextTable from the first invoke")
+	}
+
+	// 2. Second invoke with an UNKNOWN payload_type (99). The Rust
+	//    side validates payload_type against its registry; an
+	//    unknown value forces a RC_PARSE_ERROR return. Critically,
+	//    that error path is now POST-consumption of the input
+	//    context_table.
+	const unknownPayloadType uint8 = 99
+	_, _, _, err = mgr.InvokeByName("hook2", unknownPayloadType, payload, &Extensions{}, ctxTable)
+	if err == nil {
+		t.Fatal("expected error from invoke with unknown payload_type")
+	}
+
+	// 3. Per the P0-1 contract, ctxTable's handle was nil'd in Go
+	//    *before* the C call returned. So whether or not Rust wrote
+	//    *context_table_out, our local handle is nil.
+	if ctxTable.handle != nil {
+		t.Errorf("input ContextTable.handle should be nil after invoke error; got %p", ctxTable.handle)
+	}
+
+	// 4. The actual UAF check: Close() must be safe. Pre-fix this
+	//    called cpex_release_context_table on already-freed memory.
+	//    Post-fix, Close() short-circuits on a nil handle and is a
+	//    no-op. Either it crashes (fail) or it doesn't (pass).
+	ctxTable.Close()
+}
+
+// TestInvokeByNameConsumesContextTableOnRcError pins the other half
+// of the P0-1 contract — even when the manager rejects the call with
+// a validation-class error (here: shutdown after first invoke), the
+// caller's ContextTable handle is nil'd unconditionally.
+//
+// Verifies: no leak of the input Box when Rust never gets to write
+// the output; Close() on the input is a safe no-op.
+func TestInvokeByNameConsumesContextTableEvenOnShutdownPath(t *testing.T) {
+	mgr, err := NewPluginManagerDefault()
+	if err != nil {
+		t.Fatalf("NewPluginManagerDefault failed: %v", err)
+	}
+	if err := mgr.Initialize(); err != nil {
+		mgr.Shutdown()
+		t.Fatalf("Initialize failed: %v", err)
+	}
+
+	payload := map[string]any{"tool_name": "test"}
+	_, ctxTable, bg, err := mgr.InvokeByName("hook1", PayloadGeneric, payload, &Extensions{}, nil)
+	if err != nil {
+		mgr.Shutdown()
+		t.Fatalf("first invoke failed: %v", err)
+	}
+	bg.Close()
+
+	// Shut down the manager — Go-side short-circuit will return
+	// ErrCpexInvalidHandle WITHOUT calling cpex_invoke. The Go
+	// wrapper hasn't touched ctxTable yet in this case (early
+	// return at m.handle == nil), so ctxTable.handle remains live.
+	mgr.Shutdown()
+
+	_, _, _, err = mgr.InvokeByName("hook2", PayloadGeneric, payload, &Extensions{}, ctxTable)
+	if !errors.Is(err, ErrCpexInvalidHandle) {
+		t.Errorf("expected ErrCpexInvalidHandle after shutdown, got %v", err)
+	}
+
+	// Even though the Go-side short-circuit didn't transit our
+	// handle to Rust, Close() must still be safe — it's a legal
+	// thing for callers to do.
+	ctxTable.Close()
+}