gemini-code-assist[bot] commented on code in PR #658:
URL: https://github.com/apache/tvm-ffi/pull/658#discussion_r3539075584
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -60,6 +66,122 @@ 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) {}
+
+ // All bounds checks use subtraction (size_ - cursor_ never underflows
because
+ // cursor_ <= size_ is an invariant) to avoid cursor_ + n overflowing.
+ uint64_t ReadU64() {
+ TVM_FFI_CHECK(size_ - cursor_ >= sizeof(uint64_t), 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(nbytes <= size_ - cursor_, 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();
+ // Guard reserve() against a corrupt count: each element needs 8 more
bytes,
+ // so count can't exceed the remaining byte budget. Prevents an OOM abort.
+ TVM_FFI_CHECK(count <= (size_ - cursor_) / sizeof(uint64_t), RuntimeError)
+ << "Corrupt library binary: vector size exceeds remaining blob";
+ 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";
Review Comment:
If `indptr` has fewer than 2 elements (e.g., `indptr.size() == 1`),
`num_modules` will be calculated as `0`. This leads to an empty `modules`
vector, and the subsequent `return modules[0];` on line 183 will result in an
out-of-bounds access (undefined behavior). Checking `indptr.size() >= 2`
ensures that `num_modules` is at least 1, making the return statement safe.
```c
TVM_FFI_CHECK(indptr.size() >= 2, RuntimeError)\n << \"Corrupt
library binary: import tree indptr must have at least 2 elements\";
```
--
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]
