Vulnerable code sample

import uuid
from typing import Dict, Any

# This vulnerable implementation uses a simple global dictionary for task storage.
# The crucial flaw is that it does not associate a task with the session
# that created it, allowing any session to interact with any task.
TASK_STORAGE: Dict[str, Dict[str, Any]] = {}

class VulnerableMCPServer:
    """
    A simplified representation of the server behavior before the fix.
    The handlers for task operations lack session-based authorization.
    """

    def __init__(self):
        # In a real scenario, this would be a long-lived object.
        # It operates on the global TASK_STORAGE.
        self.tasks = TASK_STORAGE

    def experimental_enable_tasks(self):
        """
        Conceptually enables the vulnerable task handlers. In the real SDK,
        this would register handlers for specific network endpoints.
        """
        # The methods below simulate the logic of the vulnerable handlers.
        print("Vulnerable task handlers for list, get, result, and cancel are enabled.")

    def create_task(self, session_id: str, data: Any) -> str:
        """
        Simulates a client from a specific session creating a task.
        """
        task_id = f"task_{uuid.uuid4().hex[:8]}"
        
        # FLAW: The creating session_id is received but NOT stored with the task.
        self.tasks[task_id] = {
            "status": "running",
            "result": None,
            "data": data,
        }
        return task_id

    # --- VULNERABLE HANDLER IMPLEMENTATIONS ---
    # The following methods simulate the default handlers. They only use 'task_id'
    # and have no concept of which session is making the request.

    def handle_list_tasks(self) -> list:
        """
        VULNERABLE: Any client can list all task IDs on the server.
        This handler does not filter by session.
        """
        return list(self.tasks.keys())

    def handle_get_task(self, task_id: str) -> Dict[str, Any]:
        """
        VULNERABLE: Any client can get the details of any task if they
        know or can guess the task_id.
        """
        if task_id not in self.tasks:
            raise KeyError("Task not found")
        return self.tasks[task_id]

    def handle_get_result(self, task_id: str) -> Any:
        """
        VULNERABLE: Any client can retrieve the result of any task.
        """
        if task_id not in self.tasks:
            raise KeyError("Task not found")
        # Simulate task completion for demonstration
        if self.tasks[task_id].get("result") is None:
             self.tasks[task_id]["result"] = f"Sensitive result for {task_id}"
        return self.tasks[task_id]["result"]

    def handle_cancel_task(self, task_id: str) -> Dict[str, str]:
        """
        VULNERABLE: Any client can cancel any other client's task.
        """
        if task_id not in self.tasks:
            raise KeyError("Task not found")
        
        task = self.tasks[task_id]
        if task["status"] == "running":
            task["status"] = "cancelled"
            return {"status": "success", "message": f"Task {task_id} has been cancelled."}
        return {"status": "noop", "message": "Task was not in a running state."}