gemini-code-assist[bot] commented on code in PR #658:
URL: https://github.com/apache/tvm-ffi/pull/658#discussion_r3539310139
##########
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));
Review Comment:
The `nbytes` header read from the embedded library binary is untrusted and
can be arbitrarily large (e.g., due to corruption or malicious input). Relying
on it directly to size the `BlobReader` and subsequently guard vector/string
allocations can lead to integer overflow bypasses and process crashes via
out-of-memory (OOM) or allocation failures (e.g., in `reserve()`). Enforcing a
sane upper limit on `nbytes` (such as 2 GB) prevents these issues.
```c
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);
}
// Sanity check to prevent integer overflow or absurd memory allocation if
the header is corrupt.
constexpr uint64_t kMaxLibraryBinSize = 2ULL * 1024 * 1024 * 1024; // 2 GB
TVM_FFI_CHECK(nbytes < kMaxLibraryBinSize, RuntimeError)
<< "Corrupt library binary: size header " << nbytes << " exceeds
maximum allowed size (2 GB)";
BlobReader reader(library_bin + sizeof(nbytes),
static_cast<size_t>(nbytes));
```
--
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]
