fix: add safety warning to pickle_memoize cache loading#18925
fix: add safety warning to pickle_memoize cache loading#18925tlopex merged 1 commit intoapache:mainfrom
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a security vulnerability (CWE-502: Deserialization of Untrusted Data) in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a UserWarning to alert users about the security risks of loading cached pickle files. This is a good security measure. I have two suggestions to improve this change: first, to move the import statement to the top of the file to follow Python best practices. Second, to add a test to verify that the warning is triggered as expected. Please see the detailed comments on the code.
python/tvm/contrib/pickle_memoize.py
Outdated
| if self.path.exists(): | ||
| with self.path.open("rb") as cache_file: | ||
| try: | ||
| import warnings |
There was a problem hiding this comment.
To adhere to PEP 8 guidelines, import statements should be placed at the top of the file. Please move import warnings to the top of the module.
References
- PEP 8, the style guide for Python code, specifies that imports should be at the top of the file. This improves readability and makes dependencies explicit. (link)
| warnings.warn( | ||
| f"Loading cached pickle file from {self.path}. " | ||
| "Pickle files can execute arbitrary code. " | ||
| "Only load cache files you trust.", | ||
| UserWarning, | ||
| stacklevel=2, | ||
| ) |
There was a problem hiding this comment.
This change introduces a user warning, which is great for security awareness. However, there is no test to verify that this warning is actually triggered when loading from a cached pickle file. Please consider adding a test case to tests/python/contrib/test_memoize.py that uses pytest.warns to assert that the UserWarning is raised.
There was a problem hiding this comment.
Code Review
This pull request addresses a security concern with pickle.load by adding a UserWarning to inform users of the potential for arbitrary code execution when loading cached files. The change is straightforward and effectively communicates the risk. I have one suggestion regarding the placement of the import warnings statement to improve code style and maintainability.
python/tvm/contrib/pickle_memoize.py
Outdated
| if self.path.exists(): | ||
| with self.path.open("rb") as cache_file: | ||
| try: | ||
| import warnings |
There was a problem hiding this comment.
The import warnings statement should be moved to the top of the file. According to PEP 8, imports should be at the top of the file, just after any module comments and docstrings. Placing imports inside functions is generally discouraged as it makes it harder to see the module's dependencies at a glance and can lead to repeated imports.
References
- PEP 8 recommends that imports are always put at the top of the file, just after any module comments and docstrings, and before module globals and constants. (link)
pickle_memoize loads cached pickle files via pickle.load() without any integrity verification. If an attacker can write to the cache directory, they can inject malicious pickle payloads that execute on next load. - Add UserWarning when loading pickle cache files - Move import to module level (PEP 8) - Add test case to verify warning is raised Fixes: Unsafe pickle.load() in cache (CWE-502) Signed-off-by: scruge1 <scruge1@proton.me>
scruge1
force-pushed
the
fix/pickle-memoize-unsafe-load
branch
from
f1bb09f to
728fed0
Compare
|
given pickle is used as part of our code and our security model assumes trusted data for pickle memoize that is only used for trusted data in testcase, we would prefer not having a warning in such case (just like pickle itself won't have warnings). |
…)" This reverts commit 7b3fa38.
|
Understood — makes sense that pickle_memoize operates within TVM's trusted-data security model. Closing this PR. Thank you for the review. |
|
Understood — the revert in #18926 is appropriate given TVM's trusted-data security model. Thank you for the consideration. |
3 participants
Summary
pickle_memoizeloads cached pickle files viapickle.load()without any integrity verification or user warning. If an attacker can write to the cache directory, they can inject malicious pickle payloads that execute arbitrary code on next load.Fix
Adds a
UserWarningwhen loading pickle cache files to alert users about the security risk.Related
Huntr security vulnerability report (CWE-502: Deserialization of Untrusted Data)