Vulnerable code sample

import base64
import dill
from dataclasses import dataclass

# Simplified representation of Feast's data structures involved in the vulnerability.
# In the actual Feast code, these would be Protobuf-generated classes.

@dataclass
class UserDefinedFunction:
    """Represents the user-defined function spec."""
    name: str
    body: bytes  # This field contains the base64-encoded, dill-serialized object.

@dataclass
class OnDemandFeatureView:
    """Represents the on-demand feature view spec."""
    name: str
    user_defined_function: UserDefinedFunction


def apply_on_demand_feature_view(spec: OnDemandFeatureView):
    """
    A simplified function simulating the vulnerable part of the gRPC handler
    in the Feast registry server.

    This function is called to apply a new feature view specification.
    Crucially, it deserializes user-provided data before any authorization checks.
    """
    print(f"Applying OnDemandFeatureView spec: {spec.name}")

    # --- VULNERABLE CODE BLOCK ---
    # The 'body' is expected to be a base64-encoded representation of a
    # dill-serialized Python function.
    # No validation or authorization has occurred at this point.
    if spec.user_defined_function and spec.user_defined_function.body:
        try:
            # 1. The user-controlled 'body' is decoded from base64.
            decoded_body = base64.b64decode(spec.user_defined_function.body)

            # 2. The result is passed directly to dill.loads().
            # An attacker can provide a serialized object with a malicious
            # __reduce__ method, which will be executed here, leading to RCE.
            udf = dill.loads(decoded_body)
            print("Successfully deserialized UDF.")

            # In the real application, further processing of the 'udf' object would follow.
            # However, the remote code execution has already happened during 'dill.loads()'.

        except Exception as e:
            print(f"An error occurred during deserialization: {e}")
            # The vulnerability exists even if an exception is caught,
            # as the malicious code in __reduce__ executes before the exception.

    # In a real application, authorization checks and registry updates would happen here,
    # but it's too late as the code has already been executed.
    print("Proceeding with (post-exploitation) registry operations...")