Add hash check for pickle files (CVE-2024-24590)

allegroai · Feb 9, 2024 · e506831 · e506831
1 parent b78e9b9
commit e506831
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/clearml/binding/artifacts.py b/clearml/binding/artifacts.py
@@ -203,6 +203,10 @@ def get(self, force_download=False, deserialization_function=None):
                 with open(local_file, "rt") as f:
                     self._object = f.read()
             elif self.type == "pickle":
+                if self.hash:
+                    file_hash, _ = sha256sum(local_file, block_size=Artifacts._hash_block_size)
+                    if self.hash != file_hash:
+                        raise Exception("incorrect pickle file hash, artifact file might be corrupted")
                 with open(local_file, "rb") as f:
                     self._object = pickle.load(f)
         except Exception as e: