Patched code sample

# NOTE: The provided CVE-2026-21874 is not a real CVE identifier.
# The vulnerability described matches CVE-2024-21898 for NiceGUI.
# This code provides a conceptual, self-contained demonstration of how
# that vulnerability was fixed. The fix ensures that a Redis connection,
# created for a specific user session (client), is explicitly closed when
# that client disconnects.

import threading
import time
from typing import Dict

# --- Mock Objects for Demonstration ---
# This allows the script to run without a real Redis instance.

class MockRedisConnection:
    """A mock Redis connection to simulate connection management."""
    def __init__(self, **kwargs):
        self.is_closed = False
        self.id = id(self)
        print(f"[Redis Mock] CREATED connection {self.id}")

    def close(self):
        """Simulates closing the connection."""
        self.is_closed = True
        print(f"[Redis Mock] CLOSED connection {self.id}")

    def get(self, key):
        """Simulates a Redis GET operation."""
        if self.is_closed:
            raise ConnectionError("Connection is closed")
        return f"value_for_{key}"

# --- Conceptual Representation of the Fix in NiceGUI's Storage ---

class FixedRedisStorage:
    """
    A conceptual representation of NiceGUI's fixed Redis storage.

    The key to the fix is managing connections on a per-client basis
    instead of using a single, long-lived connection or creating new ones
    without a cleanup mechanism.
    """
    def __init__(self):
        # A dictionary to hold active Redis connections, keyed by a client's unique ID.
        self._redis_connections: Dict[str, MockRedisConnection] = {}
        # The configuration for creating new connections.
        self.redis_config = {'host': 'localhost', 'port': 6379}
        print("Initialized FixedRedisStorage.")

    def _get_connection(self, client_id: str) -> MockRedisConnection:
        """
        Retrieves or creates a Redis connection for a specific client.
        This is the entry point for a client to interact with storage.
        """
        # If a connection for this client doesn't exist, create and store it.
        if client_id not in self._redis_connections:
            print(f"Client '{client_id}': No connection found. Creating a new one.")
            self._redis_connections[client_id] = MockRedisConnection(**self.redis_config)
        else:
            print(f"Client '{client_id}': Reusing existing connection.")
        return self._redis_connections[client_id]

    def _close_connection(self, client_id: str) -> None:
        """
        THE FIX: This method is called upon client disconnection.
        It closes the specific Redis connection associated with the client
        and removes it from the pool of active connections.
        """
        if client_id in self._redis_connections:
            print(f"Client '{client_id}': Disconnecting. Closing its Redis connection.")
            connection = self._redis_connections.pop(client_id)
            connection.close()
        else:
            # This case can happen but is not part of the vulnerability flow.
            print(f"Client '{client_id}': No connection found to close.")

    @property
    def active_connections(self) -> int:
        """Helper property to see the number of open connections."""
        return len(self._redis_connections)

# --- Simulation of Client Connections and Disconnections ---

class Client:
    """A mock client representing a user's browser tab."""
    def __init__(self, client_id: str, storage: FixedRedisStorage):
        self.id = client_id
        self.storage = storage
        print(f"\n---> Simulating connect for Client '{self.id}'")
        # On connect, the client requests its storage, which triggers connection creation.
        self.redis = self.storage._get_connection(self.id)
        # Simulate doing some work
        self.redis.get(f'user:{self.id}:data')

    def disconnect(self):
        """Simulates the user closing the browser tab."""
        print(f"\n---> Simulating disconnect for Client '{self.id}'")
        # On disconnect, the crucial cleanup function is called.
        self.storage._close_connection(self.id)


if __name__ == "__main__":
    # 1. The central storage object is created for the application.
    storage = FixedRedisStorage()
    print(f"Initial state: {storage.active_connections} active connections.")
    print("-" * 50)

    # 2. A user opens the app in a browser tab. A new client is created.
    client1 = Client("session_abc_123", storage)
    print(f"After Client 1 connect: {storage.active_connections} active connections.")
    print("-" * 50)

    # 3. A second user (or the same user in a new tab) connects.
    client2 = Client("session_def_456", storage)
    print(f"After Client 2 connect: {storage.active_connections} active connections.")
    print("-" * 50)

    # 4. The first user closes their browser tab. The client disconnects.
    #    The vulnerable version would NOT clean up the connection here.
    #    The FIXED version correctly closes the connection.
    client1.disconnect()
    print(f"After Client 1 disconnect: {storage.active_connections} active connections.")
    print("-" * 50)

    # 5. The second user also disconnects.
    client2.disconnect()
    print(f"After Client 2 disconnect: {storage.active_connections} active connections.")
    print("-" * 50)

    print("\nDemonstration complete.")
    print("The number of active connections correctly returned to 0, preventing a connection leak.")