CustomOp Lite example leads to use-after-free

[lunixbochs]

Describe the issue

Starting from a9df3ae

Using the OrtLiteCustomOp example from https://onnxruntime.ai/docs/reference/operators/add-custom-op.html leads to a use-after-free if your session lasts longer than the scope used to set it up.

Ort::CustomOpDomain v1_domain{"v1"};
std::unique_ptr<OrtLiteCustomOp> custom_op_one{Ort::Custom::CreateLiteCustomOp("CustomOpOne", "CPUExecutionProvider", KernelOne)};
v1_domain.Add(custom_op_one.get());
Ort::SessionOptions session_options;
session_options.Add(v1_domain);
// create a session with the session_options ...

v1_domain.Add() takes an unsafe interior pointer from the std::unique_ptr using custom_op_one.get(). If the unique_ptr falls out of scope and is deallocated while the ort session is still alive, calling the op will result in a use-after-free and crash onnxruntime.

It's trivial to encounter this if you construct an ort session in a wrapper class and reuse it repeatedly by calling a method on the class - the op will fall out of scope immediately and crash. IMO the domain should take ownership or a copy of the custom op if possible, the caller shouldn't be trusted with an implicit contract to keep the op around for longer than the session.

The tests don't encounter this because they use the ort session in the same scope as the unique_ptr.
I wasted a couple days debugging this.

To reproduce

Follow the custom op lite example, but put it in a scope block so the unique_ptr gets dropped immediately:

{
    std::unique_ptr<OrtLiteCustomOp> custom_op_one{Ort::Custom::CreateLiteCustomOp("CustomOpOne", "CPUExecutionProvider", KernelOne)};
    v1_domain.Add(custom_op_one.get());
}

Urgency

Ideally this is evaluated before 1.16.0 is actually released with custom ops, in case a breaking api change is needed.

Platform

Mac

OS Version

13.5

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

a9df3ae

ONNX Runtime API

C++

Architecture

ARM64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[pranavsharma]

@RandySheriff - can you take a look?

[RandySheriffH]

@RandySheriff - can you take a look?

If understand it correctly - the management of the lifetime of custom ops could be in a way that is appropriate to custom application.
For example, the unique ptr hosting the custom op should PROPERLY saved somewhere to coexist with the session that is consuming the op.

[lunixbochs]

I think the Ort::CustomOpDomain object is affected in the same way as it's passed as a pointer to session_options.Add

[pranavsharma]

The docs do say this.

C API:

onnxruntime/include/onnxruntime/core/session/onnxruntime_c_api.h

Line 1043 in 6c39641

* \note The OrtCustomOp* pointer must remain valid until the ::OrtCustomOpDomain using it is released

and

C++ API:

onnxruntime/include/onnxruntime/core/session/onnxruntime_cxx_api.h

Line 735 in 6c39641

// This does not take ownership of the op, simply registers it.

Lite Custom ops API (PR recently merged):
https://github.com/microsoft/onnxruntime/pull/17576/files

[lunixbochs]

If you're considering this documented, then I strongly urge you to update the example. I followed the linked example and encountered this. (But if it's reasonable for you to fix the unsafety, even better imo, as the resulting crash is nondeterministic, deep in onnxruntime, and not easily debuggable even with symbols)