gemini-code-assist[bot] commented on code in PR #658:
URL: https://github.com/apache/tvm-ffi/pull/658#discussion_r3536928369
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +65,102 @@ struct DylibFnContextWithModule {
};
void DeleteDylibFnContextWithModule(void* p) { delete
static_cast<DylibFnContextWithModule*>(p); }
+
+// Minimal little-endian reader for the embedded library-binary blob. The addon
+// links only against public tvm-ffi headers, so it re-implements the reduced
+// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All
+// supported targets are little-endian (x86_64, aarch64, arm64).
+class BlobReader {
+ public:
+ BlobReader(const char* data, size_t size) : data_(data), size_(size) {}
+
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(cursor_ + sizeof(uint64_t) <= size_, RuntimeError)
+ << "Corrupt library binary: unexpected end of blob";
Review Comment:
The check `cursor_ + sizeof(uint64_t) <= size_` is susceptible to integer
overflow if `size_` is extremely large. To prevent integer overflow, rewrite
the check using subtraction: `size_ - cursor_ >= sizeof(uint64_t)`.
```c
uint64_t ReadU64() {
TVM_FFI_CHECK(size_ - cursor_ >= sizeof(uint64_t), RuntimeError)
<< "Corrupt library binary: unexpected end of blob";
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +65,102 @@ struct DylibFnContextWithModule {
};
void DeleteDylibFnContextWithModule(void* p) { delete
static_cast<DylibFnContextWithModule*>(p); }
+
+// Minimal little-endian reader for the embedded library-binary blob. The addon
+// links only against public tvm-ffi headers, so it re-implements the reduced
+// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All
+// supported targets are little-endian (x86_64, aarch64, arm64).
+class BlobReader {
+ public:
+ BlobReader(const char* data, size_t size) : data_(data), size_(size) {}
+
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(cursor_ + sizeof(uint64_t) <= size_, RuntimeError)
+ << "Corrupt library binary: unexpected end of blob";
+ uint64_t value = 0;
+ for (size_t i = 0; i < sizeof(uint64_t); ++i) {
+ value |= static_cast<uint64_t>(static_cast<unsigned char>(data_[cursor_
+ i])) << (i * 8);
+ }
+ cursor_ += sizeof(uint64_t);
+ return value;
+ }
+
+ std::string ReadString() {
+ uint64_t nbytes = ReadU64();
+ TVM_FFI_CHECK(cursor_ + nbytes <= size_, RuntimeError)
+ << "Corrupt library binary: string length exceeds blob";
Review Comment:
The check `cursor_ + nbytes <= size_` is susceptible to integer overflow if
`nbytes` is extremely large (e.g., `0xFFFFFFFFFFFFFFFF`). If an overflow
occurs, the check can be bypassed, leading to out-of-bounds reads or excessive
memory allocation. To prevent integer overflow, rewrite the check using
subtraction: `nbytes <= size_ - cursor_`.
```c
std::string ReadString() {
uint64_t nbytes = ReadU64();
TVM_FFI_CHECK(nbytes <= size_ - cursor_, RuntimeError)
<< "Corrupt library binary: string length exceeds blob";
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +65,102 @@ struct DylibFnContextWithModule {
};
void DeleteDylibFnContextWithModule(void* p) { delete
static_cast<DylibFnContextWithModule*>(p); }
+
+// Minimal little-endian reader for the embedded library-binary blob. The addon
+// links only against public tvm-ffi headers, so it re-implements the reduced
+// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All
+// supported targets are little-endian (x86_64, aarch64, arm64).
+class BlobReader {
+ public:
+ BlobReader(const char* data, size_t size) : data_(data), size_(size) {}
+
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(cursor_ + sizeof(uint64_t) <= size_, RuntimeError)
+ << "Corrupt library binary: unexpected end of blob";
+ uint64_t value = 0;
+ for (size_t i = 0; i < sizeof(uint64_t); ++i) {
+ value |= static_cast<uint64_t>(static_cast<unsigned char>(data_[cursor_
+ i])) << (i * 8);
+ }
+ cursor_ += sizeof(uint64_t);
+ return value;
+ }
+
+ std::string ReadString() {
+ uint64_t nbytes = ReadU64();
+ TVM_FFI_CHECK(cursor_ + nbytes <= size_, RuntimeError)
+ << "Corrupt library binary: string length exceeds blob";
+ std::string out(data_ + cursor_, static_cast<size_t>(nbytes));
+ cursor_ += static_cast<size_t>(nbytes);
+ return out;
+ }
+
+ std::vector<uint64_t> ReadU64Vector() {
+ uint64_t count = ReadU64();
+ std::vector<uint64_t> out;
+ out.reserve(static_cast<size_t>(count));
+ for (uint64_t i = 0; i < count; ++i) {
+ out.push_back(ReadU64());
+ }
+ return out;
+ }
Review Comment:
In `ReadU64Vector()`, `out.reserve(static_cast<size_t>(count))` is called
with a `count` read directly from the binary blob. If `count` is extremely
large or corrupt, this will attempt to allocate a massive amount of memory,
leading to an out-of-memory crash (DoS).
Since each element in the vector is a `uint64_t` (8 bytes), the remaining
size in the blob must be at least `count * sizeof(uint64_t)`. We should
validate this before reserving memory to prevent excessive allocation and
out-of-bounds reads.
```c
std::vector<uint64_t> ReadU64Vector() {
uint64_t count = ReadU64();
TVM_FFI_CHECK(count <= (size_ - cursor_) / sizeof(uint64_t),
RuntimeError)
<< "Corrupt library binary: vector size exceeds remaining blob size";
std::vector<uint64_t> out;
out.reserve(static_cast<size_t>(count));
for (uint64_t i = 0; i < count; ++i) {
out.push_back(ReadU64());
}
return out;
}
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +65,102 @@ struct DylibFnContextWithModule {
};
void DeleteDylibFnContextWithModule(void* p) { delete
static_cast<DylibFnContextWithModule*>(p); }
+
+// Minimal little-endian reader for the embedded library-binary blob. The addon
+// links only against public tvm-ffi headers, so it re-implements the reduced
+// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All
+// supported targets are little-endian (x86_64, aarch64, arm64).
+class BlobReader {
+ public:
+ BlobReader(const char* data, size_t size) : data_(data), size_(size) {}
+
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(cursor_ + sizeof(uint64_t) <= size_, RuntimeError)
+ << "Corrupt library binary: unexpected end of blob";
+ uint64_t value = 0;
+ for (size_t i = 0; i < sizeof(uint64_t); ++i) {
+ value |= static_cast<uint64_t>(static_cast<unsigned char>(data_[cursor_
+ i])) << (i * 8);
+ }
+ cursor_ += sizeof(uint64_t);
+ return value;
+ }
+
+ std::string ReadString() {
+ uint64_t nbytes = ReadU64();
+ TVM_FFI_CHECK(cursor_ + nbytes <= size_, RuntimeError)
+ << "Corrupt library binary: string length exceeds blob";
+ std::string out(data_ + cursor_, static_cast<size_t>(nbytes));
+ cursor_ += static_cast<size_t>(nbytes);
+ return out;
+ }
+
+ std::vector<uint64_t> ReadU64Vector() {
+ uint64_t count = ReadU64();
+ std::vector<uint64_t> out;
+ out.reserve(static_cast<size_t>(count));
+ for (uint64_t i = 0; i < count; ++i) {
+ out.push_back(ReadU64());
+ }
+ return out;
+ }
+
+ private:
+ const char* data_;
+ size_t size_;
+ size_t cursor_{0};
+};
+
+// Reduced re-implementation of the core ProcessLibraryBin parser (Option A).
+// Deserializes the modules embedded in a __tvm_ffi__library_bin blob, wires up
+// the import tree, and returns the root module. The "_lib" placeholder is
+// filled by `lib_module` (the JIT dylib module itself). Custom modules are
+// deserialized through the public ffi.Module.load_from_bytes.<kind> registry.
+//
+// Blob layout (little-endian):
+// <nbytes: u64> <indptr: vec<u64>> <child_indices: vec<u64>>
+// <kind0: str> [<bytes0: str>] <kind1: str> [<bytes1: str>] ...
+// where vec<u64> = <count: u64> <u64 * count>, str = <len: u64> <bytes>,
+// and the import tree is a CSR: module i imports child_indices[indptr[i]..].
+Module ProcessEmbeddedLibraryBin(const char* library_bin, const Module&
lib_module) {
+ uint64_t nbytes = 0;
+ for (size_t i = 0; i < sizeof(nbytes); ++i) {
+ nbytes |= static_cast<uint64_t>(static_cast<unsigned
char>(library_bin[i])) << (i * 8);
+ }
+ BlobReader reader(library_bin + sizeof(nbytes), static_cast<size_t>(nbytes));
+
+ std::vector<uint64_t> indptr = reader.ReadU64Vector();
+ std::vector<uint64_t> child_indices = reader.ReadU64Vector();
+ TVM_FFI_CHECK(!indptr.empty(), RuntimeError) << "Corrupt library binary:
empty import tree";
+ size_t num_modules = indptr.size() - 1;
Review Comment:
The CSR (Compressed Sparse Row) import tree is parsed without validating
that `indptr` elements are monotonically non-decreasing, that `indptr[0] == 0`,
or that `indptr.back() <= child_indices.size()`. If any of these properties are
violated, the loop `for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j)` can
perform out-of-bounds reads on `child_indices`.
Add validation checks for the CSR structure to ensure safety.
```c
std::vector<uint64_t> indptr = reader.ReadU64Vector();
std::vector<uint64_t> child_indices = reader.ReadU64Vector();
TVM_FFI_CHECK(!indptr.empty(), RuntimeError) << "Corrupt library binary:
empty import tree";
TVM_FFI_CHECK(indptr[0] == 0, RuntimeError)
<< "Corrupt library binary: indptr[0] must be 0";
size_t num_modules = indptr.size() - 1;
for (size_t i = 0; i < num_modules; ++i) {
TVM_FFI_CHECK(indptr[i] <= indptr[i + 1], RuntimeError)
<< "Corrupt library binary: indptr must be non-decreasing";
}
TVM_FFI_CHECK(indptr.back() <= child_indices.size(), RuntimeError)
<< "Corrupt library binary: indptr.back() exceeds child_indices size";
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +65,102 @@ struct DylibFnContextWithModule {
};
void DeleteDylibFnContextWithModule(void* p) { delete
static_cast<DylibFnContextWithModule*>(p); }
+
+// Minimal little-endian reader for the embedded library-binary blob. The addon
+// links only against public tvm-ffi headers, so it re-implements the reduced
+// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All
+// supported targets are little-endian (x86_64, aarch64, arm64).
+class BlobReader {
+ public:
+ BlobReader(const char* data, size_t size) : data_(data), size_(size) {}
+
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(cursor_ + sizeof(uint64_t) <= size_, RuntimeError)
+ << "Corrupt library binary: unexpected end of blob";
+ uint64_t value = 0;
+ for (size_t i = 0; i < sizeof(uint64_t); ++i) {
+ value |= static_cast<uint64_t>(static_cast<unsigned char>(data_[cursor_
+ i])) << (i * 8);
+ }
+ cursor_ += sizeof(uint64_t);
+ return value;
+ }
+
+ std::string ReadString() {
+ uint64_t nbytes = ReadU64();
+ TVM_FFI_CHECK(cursor_ + nbytes <= size_, RuntimeError)
+ << "Corrupt library binary: string length exceeds blob";
+ std::string out(data_ + cursor_, static_cast<size_t>(nbytes));
+ cursor_ += static_cast<size_t>(nbytes);
+ return out;
+ }
+
+ std::vector<uint64_t> ReadU64Vector() {
+ uint64_t count = ReadU64();
+ std::vector<uint64_t> out;
+ out.reserve(static_cast<size_t>(count));
+ for (uint64_t i = 0; i < count; ++i) {
+ out.push_back(ReadU64());
+ }
+ return out;
+ }
+
+ private:
+ const char* data_;
+ size_t size_;
+ size_t cursor_{0};
+};
+
+// Reduced re-implementation of the core ProcessLibraryBin parser (Option A).
+// Deserializes the modules embedded in a __tvm_ffi__library_bin blob, wires up
+// the import tree, and returns the root module. The "_lib" placeholder is
+// filled by `lib_module` (the JIT dylib module itself). Custom modules are
+// deserialized through the public ffi.Module.load_from_bytes.<kind> registry.
+//
+// Blob layout (little-endian):
+// <nbytes: u64> <indptr: vec<u64>> <child_indices: vec<u64>>
+// <kind0: str> [<bytes0: str>] <kind1: str> [<bytes1: str>] ...
+// where vec<u64> = <count: u64> <u64 * count>, str = <len: u64> <bytes>,
+// and the import tree is a CSR: module i imports child_indices[indptr[i]..].
+Module ProcessEmbeddedLibraryBin(const char* library_bin, const Module&
lib_module) {
+ uint64_t nbytes = 0;
+ for (size_t i = 0; i < sizeof(nbytes); ++i) {
+ nbytes |= static_cast<uint64_t>(static_cast<unsigned
char>(library_bin[i])) << (i * 8);
+ }
+ BlobReader reader(library_bin + sizeof(nbytes), static_cast<size_t>(nbytes));
+
+ std::vector<uint64_t> indptr = reader.ReadU64Vector();
+ std::vector<uint64_t> child_indices = reader.ReadU64Vector();
+ TVM_FFI_CHECK(!indptr.empty(), RuntimeError) << "Corrupt library binary:
empty import tree";
+ size_t num_modules = indptr.size() - 1;
+
+ std::vector<Module> modules;
+ modules.reserve(num_modules);
+ for (size_t i = 0; i < num_modules; ++i) {
+ std::string kind = reader.ReadString();
+ if (kind == "_lib") {
+ // Placeholder for the symbol source — the JIT dylib module itself.
+ modules.push_back(lib_module);
+ } else {
+ std::string module_bytes = reader.ReadString();
+ auto floader = Function::GetGlobal("ffi.Module.load_from_bytes." + kind);
+ TVM_FFI_CHECK(floader.has_value(), RuntimeError)
+ << "Library binary embeds a {" << kind << "} module but loader "
+ << "ffi.Module.load_from_bytes." << kind << " is not registered";
+ modules.push_back((*floader)(Bytes(module_bytes)).cast<Module>());
+ }
+ }
+
+ // Wire the import tree (CSR) using the public ModuleObj::ImportModule.
+ for (size_t i = 0; i < modules.size(); ++i) {
+ for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j) {
+ uint64_t child = child_indices[j];
+ TVM_FFI_CHECK(child < modules.size(), RuntimeError)
+ << "Corrupt library binary: child index out of range";
+ modules[i]->ImportModule(modules[static_cast<size_t>(child)]);
+ }
+ }
Review Comment:
If a custom module loader fails or returns an uninitialized/null `Module`,
calling `ImportModule` on it will cause a segmentation fault. We should
defensively check that both the parent and child modules are defined before
importing.
```c
// Wire the import tree (CSR) using the public ModuleObj::ImportModule.
for (size_t i = 0; i < modules.size(); ++i) {
for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j) {
uint64_t child = child_indices[j];
TVM_FFI_CHECK(child < modules.size(), RuntimeError)
<< "Corrupt library binary: child index out of range";
TVM_FFI_CHECK(modules[i].defined(), RuntimeError)
<< "Corrupt library binary: module at index " << i << " is null";
TVM_FFI_CHECK(modules[static_cast<size_t>(child)].defined(),
RuntimeError)
<< "Corrupt library binary: child module at index " << child << "
is null";
modules[i]->ImportModule(modules[static_cast<size_t>(child)]);
}
}
```
##########
addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py:
##########
@@ -123,10 +183,51 @@ def clear_free_slabs(self) -> int:
returned, the C++ destructor has finished and the slab's live count
reflects the drop.
- Returns:
+ Returns
+ -------
+ int
Number of slabs actually munmap'd. Returns 0 on macOS/Windows
(slab pool compiled out) or when the pool is disabled via
``slab_size=-1``.
"""
return int(_ffi_api.ExecutionSessionClearFreeSlabs(self)) # type:
ignore
+
+
+_default_session: ExecutionSession | None = None
+
+
+def default_session() -> ExecutionSession:
+ """Return the process-wide shared execution session.
+
+ A single leaked, never-destroyed session shared by all callers in the
+ process, so they share one LLVM ``ExecutionSession`` — hence process
+ symbols, the slab arena, and cross-library linking. Created on first call
+ and cached for the lifetime of the process.
+
+ The ORC runtime path is resolved in order: an embedder-registered default
+ (via ``ffi_orcjit.SetDefaultOrcRuntimePath``), then the runtime bundled
next
+ to this extension, then none. For an isolated session or a tuned arena,
+ construct an :class:`ExecutionSession` directly instead.
+
+ Returns
+ -------
+ ExecutionSession
+ The shared execution session.
+
+ Examples
+ --------
+ >>> import tvm_ffi_orcjit as oj
+ >>> session = oj.default_session()
+ >>> mod = session.load_module("dylib0.o")
+
+ """
+ global _default_session
+ if _default_session is None:
+ # Honor an embedder-registered default; otherwise fall back to the
+ # runtime bundled next to the extension. This must happen before the
+ # first DefaultSession() call, which creates the leaked singleton.
+ if not _ffi_api.GetDefaultOrcRuntimePath() and (bundled :=
_find_orc_rt_library()):
+ _ffi_api.SetDefaultOrcRuntimePath(bundled)
+ _default_session = _ffi_api.DefaultSession()
+ return _default_session
Review Comment:
In `default_session()`, the lazy initialization of `_default_session` is not
thread-safe. Since `_ffi_api.GetDefaultOrcRuntimePath()` and
`_ffi_api.SetDefaultOrcRuntimePath()` are FFI calls that may release the Python
GIL, multiple threads calling `default_session()` concurrently during the first
invocation could race, leading to multiple initializations of the process-wide
shared session.
Use a `threading.Lock` to ensure thread-safe lazy initialization.
```python
_default_session: ExecutionSession | None = None
_session_lock = __import__("threading").Lock()
def default_session() -> ExecutionSession:
"""Return the process-wide shared execution session.
A single leaked, never-destroyed session shared by all callers in the
process, so they share one LLVM ``ExecutionSession`` — hence process
symbols, the slab arena, and cross-library linking. Created on first call
and cached for the lifetime of the process.
The ORC runtime path is resolved in order: an embedder-registered default
(via ``ffi_orcjit.SetDefaultOrcRuntimePath``), then the runtime bundled
next
to this extension, then none. For an isolated session or a tuned arena,
construct an :class:`ExecutionSession` directly instead.
Returns
-------
ExecutionSession
The shared execution session.
Examples
--------
>>> import tvm_ffi_orcjit as oj
>>> session = oj.default_session()
>>> mod = session.load_module("dylib0.o")
"""
global _default_session
if _default_session is None:
with _session_lock:
if _default_session is None:
# Honor an embedder-registered default; otherwise fall back
to the
# runtime bundled next to the extension. This must happen
before the
# first DefaultSession() call, which creates the leaked
singleton.
if not _ffi_api.GetDefaultOrcRuntimePath() and (bundled :=
_find_orc_rt_library()):
_ffi_api.SetDefaultOrcRuntimePath(bundled)
_default_session = _ffi_api.DefaultSession()
return _default_session
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
