(tvm) branch main updated: fix: add safety warning to pickle_memoize cache loading (#18925)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 7b3fa38e6b fix: add safety warning to pickle_memoize cache loading 
(#18925)
7b3fa38e6b is described below

commit 7b3fa38e6b76564dabe5d6b09023a566a1520c27
Author: scruge1 <190668909+scru...@users.noreply.github.com>
AuthorDate: Tue Mar 24 05:39:47 2026 +0000

    fix: add safety warning to pickle_memoize cache loading (#18925)
    
    ## Summary
    `pickle_memoize` loads cached pickle files via `pickle.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 `UserWarning` when loading pickle cache files to alert users
    about the security risk.
    
    ## Related
    Huntr security vulnerability report (CWE-502: Deserialization of
    Untrusted Data)
    
    Signed-off-by: scruge1 <scru...@proton.me>
    Co-authored-by: scruge1 <scru...@proton.me>
---
 python/tvm/contrib/pickle_memoize.py               |  8 +++++++
 .../python/contrib/test_pickle_memoize_warning.py  | 26 ++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/python/tvm/contrib/pickle_memoize.py 
b/python/tvm/contrib/pickle_memoize.py
index f9981a5ea4..64a88068ec 100644
--- a/python/tvm/contrib/pickle_memoize.py
+++ b/python/tvm/contrib/pickle_memoize.py
@@ -23,6 +23,7 @@ import functools
 import os
 import pathlib
 import sys
+import warnings
 
 try:
     import cPickle as pickle
@@ -71,6 +72,13 @@ class Cache:
         if self.path.exists():
             with self.path.open("rb") as cache_file:
                 try:
+                    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,
+                    )
                     cache = pickle.load(cache_file)
                 except pickle.UnpicklingError:
                     cache = {}
diff --git a/tests/python/contrib/test_pickle_memoize_warning.py 
b/tests/python/contrib/test_pickle_memoize_warning.py
new file mode 100644
index 0000000000..3f26961d20
--- /dev/null
+++ b/tests/python/contrib/test_pickle_memoize_warning.py
@@ -0,0 +1,26 @@
+import pytest
+import pickle
+import tempfile
+import os
+
+
+def test_pickle_memoize_warns_on_cache_load():
+    """Test that loading a cached pickle file emits a UserWarning."""
+    from tvm.contrib.pickle_memoize import memoize
+
+    # Create a cache file
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cache_path = os.path.join(tmpdir, "test_cache")
+
+        @memoize("test_warning_cache")
+        def dummy_func():
+            return 42
+
+        # First call creates cache
+        result = dummy_func()
+        assert result == 42
+
+        # Second call loads from cache — should warn
+        with pytest.warns(UserWarning, match="Pickle files can execute 
arbitrary code"):
+            result2 = dummy_func()
+            assert result2 == 42