gemini-code-assist[bot] commented on code in PR #18662:
URL: https://github.com/apache/tvm/pull/18662#discussion_r2691399549
##########
python/tvm/ir/base.py:
##########
@@ -125,10 +125,17 @@ def load_json(json_str) -> Object:
The loaded tvm node.
"""
+ # Prevent pathological memory / CPU usage from unbounded JSON inputs.
+ MAX_JSON_BYTES = 5 * 1024 * 1024
+
+ if len(json_str) > MAX_JSON_BYTES:
+ raise ValueError("JSON IR input too large")
Review Comment:
While adding a size limit is a good security measure, a hardcoded 5MB limit
might be too restrictive for some users and could be a breaking change for
workflows involving large models.
It would be more flexible to make this limit configurable via an environment
variable, e.g., `TVM_LOAD_JSON_MAX_SIZE_BYTES`. This would provide a safe
default while allowing users to adjust it for their specific needs. A value of
`-1` could be used to disable the check entirely for trusted environments.
Additionally, constants like `MAX_JSON_BYTES` are typically defined at the
module level. If you make it configurable, you could read the environment
variable at the module level as well.
Example implementation:
```python
# At module level (requires `import os`)
_DEFAULT_MAX_JSON_BYTES = 5 * 1024 * 1024 # 5MB
_MAX_JSON_BYTES = int(os.environ.get("TVM_LOAD_JSON_MAX_SIZE_BYTES",
_DEFAULT_MAX_JSON_BYTES))
# In load_json()
if _MAX_JSON_BYTES > 0 and len(json_str) > _MAX_JSON_BYTES:
raise ValueError(
f"JSON IR input is too large ({len(json_str)} bytes >
{_MAX_JSON_BYTES} bytes). "
f"This limit is configurable via TVM_LOAD_JSON_MAX_SIZE_BYTES env
var."
)
```
--
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]
