fix: secure pickle deserialization in IPFS storage service (issue #22)
Some checks failed
AITBC CI/CD Pipeline / lint-and-test (3.11) (pull_request) Has been cancelled
AITBC CI/CD Pipeline / lint-and-test (3.12) (pull_request) Has been cancelled
AITBC CI/CD Pipeline / lint-and-test (3.13) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (apps/coordinator-api/src) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (cli/aitbc_cli) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (packages/py/aitbc-core/src) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (packages/py/aitbc-crypto/src) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (packages/py/aitbc-sdk/src) (pull_request) Has been cancelled
Security Scanning / Bandit Security Scan (tests) (pull_request) Has been cancelled
Security Scanning / CodeQL Security Analysis (javascript) (pull_request) Has been cancelled
Security Scanning / CodeQL Security Analysis (python) (pull_request) Has been cancelled
Security Scanning / Dependency Security Scan (pull_request) Has been cancelled
Security Scanning / Container Security Scan (pull_request) Has been cancelled
Security Scanning / OSSF Scorecard (pull_request) Has been cancelled
AITBC CI/CD Pipeline / test-cli (pull_request) Has been cancelled
AITBC CI/CD Pipeline / test-services (pull_request) Has been cancelled
AITBC CI/CD Pipeline / test-production-services (pull_request) Has been cancelled
AITBC CI/CD Pipeline / security-scan (pull_request) Has been cancelled
AITBC CI/CD Pipeline / build (pull_request) Has been cancelled
AITBC CI/CD Pipeline / deploy-staging (pull_request) Has been cancelled
AITBC CI/CD Pipeline / deploy-production (pull_request) Has been cancelled
AITBC CI/CD Pipeline / performance-test (pull_request) Has been cancelled
AITBC CI/CD Pipeline / docs (pull_request) Has been cancelled
AITBC CI/CD Pipeline / release (pull_request) Has been cancelled
AITBC CI/CD Pipeline / notify (pull_request) Has been cancelled
Security Scanning / Security Summary Report (pull_request) Has been cancelled

- Introduced RestrictedUnpickler in secure_pickle.py to prevent arbitrary code execution
- Updated IPFSStorageService.retrieve_memory and decompress_memory to use safe_loads()
- Maintains pickle dumps for serialization (safe)
- Reduces risk of remote code execution via malicious pickled data
This commit is contained in:
2026-03-15 21:22:51 +00:00
parent 6935c2d13d
commit 1730f3e416
2 changed files with 37 additions and 3 deletions

View File

@@ -12,6 +12,7 @@ import json
import hashlib
import gzip
import pickle
from .secure_pickle import safe_loads
from dataclasses import dataclass, asdict
try:
@@ -190,8 +191,8 @@ class IPFSStorageService:
else:
decompressed_data = retrieved_data
# Deserialize
memory_data = pickle.loads(decompressed_data)
# Deserialize (using safe unpickler)
memory_data = safe_loads(decompressed_data)
logger.info(f"Retrieved memory for agent {metadata.agent_id}: CID {cid}")
return memory_data, metadata
@@ -353,7 +354,7 @@ class MemoryCompressionService:
def decompress_memory(compressed_data: bytes) -> Any:
"""Decompress memory data"""
decompressed = gzip.decompress(compressed_data)
return pickle.loads(decompressed)
return safe_loads(decompressed)
@staticmethod
def calculate_similarity(data1: Any, data2: Any) -> float:

View File

@@ -0,0 +1,33 @@
"""
Secure pickle deserialization utilities to prevent arbitrary code execution.
"""
import pickle
import io
from typing import Any
# Safe classes whitelist: builtins and common types
SAFE_MODULES = {
'builtins': {
'list', 'dict', 'set', 'tuple', 'int', 'float', 'str', 'bytes',
'bool', 'NoneType', 'range', 'slice', 'memoryview', 'complex'
},
'datetime': {'datetime', 'date', 'time', 'timedelta', 'timezone'},
'collections': {'OrderedDict', 'defaultdict', 'Counter', 'namedtuple'},
'dataclasses': {'dataclass'},
'typing': {'Any', 'List', 'Dict', 'Tuple', 'Set', 'Optional', 'Union', 'TypeVar', 'Generic', 'NamedTuple', 'TypedDict'},
}
class RestrictedUnpickler(pickle.Unpickler):
"""
Unpickler that restricts which classes can be instantiated.
Only allows classes from SAFE_MODULES whitelist.
"""
def find_class(self, module: str, name: str) -> Any:
if module in SAFE_MODULES and name in SAFE_MODULES[module]:
return super().find_class(module, name)
raise pickle.UnpicklingError(f"Class {module}.{name} is not allowed for unpickling (security risk).")
def safe_loads(data: bytes) -> Any:
"""Safely deserialize a pickle byte stream."""
return RestrictedUnpickler(io.BytesIO(data)).load()