Patched code sample

import numpy as np
import triton_python_backend_utils as pb_utils

# This code provides a hypothetical fix for a vulnerability matching the
# description of CVE-2025-23319. The vulnerability involves an out-of-bounds
# write in the Triton Python backend, triggered by a malicious request.
#
# The scenario assumes a model that takes two inputs:
# 1. 'DATA': A numpy array of data to be written.
# 2. 'WRITE_INDEX': A scalar integer specifying where to write the data in a
#    pre-allocated buffer.
#
# The vulnerability occurs if the 'WRITE_INDEX' provided by the user is not
# validated, allowing a write operation to occur outside the buffer's bounds.

class TritonPythonModel:
    """
    A mock Triton Python backend model that demonstrates the fix for a potential
    out-of-bounds write vulnerability.
    """
    def initialize(self, args):
        """
        Initializes the model and pre-allocates a fixed-size buffer.
        """
        # In a real scenario, this buffer might be used for accumulating results.
        self.output_buffer = np.zeros(1024, dtype=np.float32)
        self.buffer_size = self.output_buffer.size

    def execute(self, requests):
        """
        This function is called for each inference request. It contains the
        vulnerability fix.
        """
        responses = []
        for request in requests:
            # 1. Retrieve user-controlled inputs from the request.
            input_data_tensor = pb_utils.get_input_tensor_by_name(request, "DATA")
            write_index_tensor = pb_utils.get_input_tensor_by_name(request, "WRITE_INDEX")

            input_data = input_data_tensor.as_numpy()
            # Ensure index is a single integer scalar.
            write_index = write_index_tensor.as_numpy().item()

            # --- VULNERABLE CODE (FOR ILLUSTRATION) ---
            #
            # Without a bounds check, an attacker could supply a `write_index`
            # that, when combined with the size of `input_data`, would write
            # past the end of `self.output_buffer`.
            #
            # Example: write_index = 1000, input_data.size = 50
            # The write would go from index 1000 to 1049, while the buffer
            # only goes up to index 1023.
            #
            # VULNERABLE_WRITE_TARGET = self.output_buffer[write_index : write_index + input_data.size]
            # VULNERABLE_WRITE_TARGET[...] = input_data
            #

            # --- FIXED CODE ---
            #
            # The fix is to rigorously validate that the entire write operation
            # will be contained within the allocated buffer's boundaries before
            # performing the write.

            # 2. Perform comprehensive bounds checking.
            is_valid = True
            error_message = ""

            # Check if start index is negative.
            if write_index < 0:
                is_valid = False
                error_message = f"Invalid WRITE_INDEX: {write_index}. Index cannot be negative."

            # Check if the write operation exceeds the buffer's upper bound.
            # This is the most critical check to prevent the out-of-bounds write.
            elif (write_index + input_data.size) > self.buffer_size:
                is_valid = False
                error_message = (
                    f"Invalid inputs would cause out-of-bounds write. "
                    f"Buffer size: {self.buffer_size}, "
                    f"requested write end: {write_index + input_data.size}."
                )

            # 3. Proceed only if the inputs are valid.
            if is_valid:
                # The write operation is now proven to be safe.
                write_slice = slice(write_index, write_index + input_data.size)
                self.output_buffer[write_slice] = input_data

                # Create a successful response.
                # Here, we return a tensor confirming the write was successful.
                out_tensor = pb_utils.Tensor(
                    "STATUS",
                    np.array(["SUCCESS"], dtype=np.object_)
                )
                response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
                responses.append(response)
            else:
                # 4. If inputs are invalid, reject the request with an error.
                # This prevents the dangerous operation and informs the client.
                error = pb_utils.TritonError(message=error_message)
                response = pb_utils.InferenceResponse(output_tensors=[], error=error)
                responses.append(response)

        return responses

    def finalize(self):
        """
        Called when the model is unloaded.
        """
        print("Finalizing the model.")