litellm/tests/integration/security/test_container_ownership.py

import json
import uuid
from collections.abc import Callable
from email.parser import BytesParser
from email.policy import HTTP
from typing import Final
from urllib.parse import parse_qsl, urlsplit

import httpx
import pytest
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, Wire, wire_server
from openai import OpenAI, PermissionDeniedError
from openai.types import ContainerListResponse, ContainerRetrieveResponse
from openai.types.containers import FileListResponse, FileRetrieveResponse
from pydantic import BaseModel, TypeAdapter

from litellm.responses.utils import ResponsesAPIRequestUtils

_FILE_ID: Final = "cfile_ownership"
_FILE_BYTES: Final = b"tenant secret contents"
_FORBIDDEN: Final = {"detail": "Forbidden"}
# Non-admins overfetch upstream pages of 100, then trim to client limit; order/after and admin limits pass unchanged.
_OWNED_LIST_UPSTREAM_PAGE_SIZE: Final = 100
_Part = tuple[str, str | None, str, bytes]


class _Deleted(BaseModel):
    id: str
    object: str
    deleted: bool


class _ContainerPage(BaseModel):
    object: str
    data: tuple[ContainerListResponse, ...]
    first_id: str | None
    last_id: str | None
    has_more: bool


_ERROR: Final = TypeAdapter(dict[str, str])


def _container(container_id: str, name: str) -> dict[str, object]:
    return {"id": container_id, "object": "container", "created_at": 1, "status": "running", "name": name}


def _file(container_id: str) -> dict[str, object]:
    return {
        "id": _FILE_ID,
        "object": "container.file",
        "container_id": container_id,
        "created_at": 1,
        "bytes": len(_FILE_BYTES),
        "path": "/mnt/data/notes.txt",
        "source": "user",
    }


def _multipart_parts(request: Request) -> tuple[_Part, ...]:
    content_type: Final = request.headers["content-type"]
    assert content_type.startswith("multipart/form-data; boundary="), content_type
    envelope: Final = f"content-type: {content_type}\r\n\r\n".encode() + request.body
    parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
    assert parsed.is_multipart(), request.body
    return tuple(
        (
            str(part.get_param("name", header="content-disposition")),
            part.get_filename(),
            part.get_content_type(),
            bytes(part.get_payload(decode=True)),
        )
        for part in parsed.iter_parts()
    )


def _container_upstream(container_ids: tuple[str, ...]) -> Callable[[Request], Reply]:
    def respond(request: Request) -> Reply:
        path: Final = urlsplit(request.target).path
        segments: Final = path.removeprefix("/v1/containers").strip("/").split("/")
        if request.method == "POST" and path == "/v1/containers":
            name: Final = str(json.loads(request.body)["name"])
            return Reply(body=json.dumps(_container(container_ids[int(name.rsplit("-", 1)[1])], name)).encode())
        container_id: Final = segments[0]
        match request.method, segments[1:]:
            case "GET", []:
                return Reply(body=json.dumps(_container(container_id, "c-0")).encode())
            case "DELETE", []:
                return Reply(
                    body=json.dumps({"id": container_id, "object": "container.deleted", "deleted": True}).encode()
                )
            case "GET", ["files"]:
                return Reply(
                    body=json.dumps(
                        {
                            "object": "list",
                            "data": [_file(container_id)],
                            "first_id": _FILE_ID,
                            "last_id": _FILE_ID,
                            "has_more": False,
                        }
                    ).encode()
                )
            case "POST", ["files"]:
                assert _multipart_parts(request) == (("file", "notes.txt", "text/plain", _FILE_BYTES),), request.body
                return Reply(body=json.dumps(_file(container_id)).encode())
            case "GET", ["files", _]:
                return Reply(body=json.dumps(_file(container_id)).encode())
            case "GET", ["files", _, "content"]:
                return Reply(body=_FILE_BYTES, content_type="application/octet-stream")
            case "DELETE", ["files", _]:
                return Reply(
                    body=json.dumps({"id": _FILE_ID, "object": "container.file.deleted", "deleted": True}).encode()
                )
        return Reply(status=404, body=b'{"error": {"message": "unscripted container route"}}')

    return respond


def _user_key(scenario: Scenario) -> tuple[str, str]:
    user_id: Final = scenario.user(user_role="internal_user")
    return user_id, scenario.key(user_id=user_id)


def _client(gateway: Gateway, key: str, model: str) -> OpenAI:
    client: Final = OpenAI(base_url=str(gateway.client.base_url.join("/v1")), api_key=key, max_retries=0)
    eventually(lambda: tuple(entry.id for entry in client.models.list()), lambda ids: model in ids)
    return client


def _routable_over_alias(gateway: Gateway, model: str) -> None:
    def listed() -> tuple[str, ...]:
        response: Final = gateway.client.get("/models", headers={"Authorization": f"Bearer {gateway.key}"})
        assert response.status_code == 200, response.text
        return tuple(str(entry["id"]) for entry in response.json()["data"])

    eventually(listed, lambda ids: model in ids)


def _create(client: OpenAI, model: str, index: int) -> str:
    created: Final = client.containers.create(name=f"c-{index}", extra_body={"model": model})
    assert created.object == "container", created.model_dump_json()
    return created.id


def _owner_rows(original_container_id: str) -> list[dict[str, object]]:
    return read_rows(
        'SELECT unified_object_id, file_purpose, created_by FROM "LiteLLM_ManagedObjectTable" WHERE model_object_id=%s',
        (f"container:openai:{original_container_id}",),
    )


def _requests(wire: Wire) -> list[tuple[str, str, str]]:
    return [(request.method, request.target, request.headers["authorization"]) for request in wire.drain()]


def _refused_with_sdk(call: Callable[[], object]) -> None:
    with pytest.raises(PermissionDeniedError) as refused:
        call()
    assert refused.value.status_code == 403, refused.value.response.text
    assert _ERROR.validate_json(refused.value.response.text) == _FORBIDDEN, refused.value.response.text


def _refused_over_alias(response: httpx.Response) -> None:
    assert response.status_code == 403, response.text
    assert _ERROR.validate_json(response.text) == _FORBIDDEN, response.text


def test_container_routes_refuse_another_users_key_and_serve_the_owner(gateway: Gateway) -> None:
    original: Final = f"cntr_{uuid.uuid4().hex}"
    provider_key: Final = f"container-key-{uuid.uuid4().hex}"
    with wire_server(_container_upstream((original,))) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=provider_key)
        owner_id, owner_key = _user_key(scenario)
        _other_id, other_key = _user_key(scenario)
        owner: Final = _client(gateway, owner_key, model)
        other: Final = _client(gateway, other_key, model)
        _routable_over_alias(gateway, model)
        created: Final = owner.containers.create(
            name="c-0",
            expires_after={"anchor": "last_active_at", "minutes": 20},
            file_ids=["file-seed"],
            extra_body={"model": model},
        )
        container_id: Final = created.id
        assert ResponsesAPIRequestUtils.decode_container_id_to_original(container_id) == original, container_id
        assert _owner_rows(original) == [
            {"unified_object_id": container_id, "file_purpose": "container", "created_by": owner_id}
        ]
        create_requests: Final = wire.drain()
        assert [
            (request.method, request.target, request.headers["authorization"]) for request in create_requests
        ] == [("POST", "/v1/containers", f"Bearer {provider_key}")]
        assert json.loads(create_requests[0].body) == {
            "name": "c-0",
            "expires_after": {"anchor": "last_active_at", "minutes": 20},
            "file_ids": ["file-seed"],
        }, create_requests[0].body

        _refused_with_sdk(lambda: other.containers.retrieve(container_id))
        _refused_with_sdk(lambda: other.containers.files.list(container_id))
        _refused_with_sdk(lambda: other.containers.files.create(container_id, file=("x.txt", b"x", "text/plain")))
        _refused_with_sdk(lambda: other.containers.files.retrieve(_FILE_ID, container_id=container_id))
        _refused_with_sdk(lambda: other.containers.files.content.retrieve(_FILE_ID, container_id=container_id))
        _refused_with_sdk(lambda: other.containers.files.delete(_FILE_ID, container_id=container_id))
        _refused_with_sdk(lambda: other.containers.delete(container_id))
        alias: Final = f"/containers/{container_id}"
        other_auth: Final = {"Authorization": f"Bearer {other_key}"}
        _refused_over_alias(gateway.client.get(alias, headers=other_auth))
        _refused_over_alias(gateway.client.get(f"{alias}/files", headers=other_auth))
        _refused_over_alias(
            gateway.client.post(f"{alias}/files", files={"file": ("x.txt", b"x", "text/plain")}, headers=other_auth)
        )
        _refused_over_alias(gateway.client.get(f"{alias}/files/{_FILE_ID}", headers=other_auth))
        _refused_over_alias(gateway.client.get(f"{alias}/files/{_FILE_ID}/content", headers=other_auth))
        _refused_over_alias(gateway.client.delete(f"{alias}/files/{_FILE_ID}", headers=other_auth))
        _refused_over_alias(gateway.client.delete(alias, headers=other_auth))
        assert _requests(wire) == []

        retrieved: Final = owner.containers.retrieve(container_id)
        assert retrieved == ContainerRetrieveResponse.model_validate(_container(original, "c-0")), retrieved
        over_alias: Final = gateway.client.get(alias, headers={"Authorization": f"Bearer {owner_key}"})
        assert over_alias.status_code == 200, over_alias.text
        assert ContainerRetrieveResponse.model_validate_json(over_alias.text) == retrieved, over_alias.text
        listed: Final = owner.containers.files.list(container_id)
        assert listed.data == [FileListResponse.model_validate(_file(original))], listed
        uploaded: Final = owner.containers.files.create(container_id, file=("notes.txt", _FILE_BYTES, "text/plain"))
        assert uploaded.id == _FILE_ID, uploaded
        fetched: Final = owner.containers.files.retrieve(_FILE_ID, container_id=container_id)
        assert fetched == FileRetrieveResponse.model_validate(_file(original)), fetched
        content: Final = owner.containers.files.content.retrieve(_FILE_ID, container_id=container_id)
        assert content.read() == _FILE_BYTES
        file_deleted: Final = gateway.client.delete(
            f"/v1/containers/{container_id}/files/{_FILE_ID}", headers={"Authorization": f"Bearer {owner_key}"}
        )
        assert file_deleted.status_code == 200, file_deleted.text
        assert _Deleted.model_validate_json(file_deleted.text) == _Deleted(
            id=_FILE_ID, object="container.file.deleted", deleted=True
        ), file_deleted.text
        deleted: Final = gateway.client.delete(
            f"/v1/containers/{container_id}", headers={"Authorization": f"Bearer {owner_key}"}
        )
        assert deleted.status_code == 200, deleted.text
        assert _Deleted.model_validate_json(deleted.text) == _Deleted(
            id=original, object="container.deleted", deleted=True
        ), deleted.text
        prefix: Final = f"/v1/containers/{original}"
        auth: Final = f"Bearer {provider_key}"
        assert _requests(wire) == [
            ("GET", prefix, auth),
            ("GET", prefix, auth),
            ("GET", f"{prefix}/files", auth),
            ("POST", f"{prefix}/files", auth),
            ("GET", f"{prefix}/files/{_FILE_ID}", auth),
            ("GET", f"{prefix}/files/{_FILE_ID}/content", auth),
            ("DELETE", f"{prefix}/files/{_FILE_ID}", auth),
            ("DELETE", prefix, auth),
        ]


def _listing_upstream(
    container_ids: tuple[str, ...], pages: dict[str | None, tuple[tuple[str, ...], bool]]
) -> Callable[[Request], Reply]:
    created: Final = _container_upstream(container_ids)

    def respond(request: Request) -> Reply:
        if request.method != "GET" or urlsplit(request.target).path != "/v1/containers":
            return created(request)
        after: Final = dict(parse_qsl(urlsplit(request.target).query)).get("after")
        ids, has_more = pages[after]
        return Reply(
            body=json.dumps(
                {
                    "object": "list",
                    "data": [_container(container_id, container_id) for container_id in ids],
                    "first_id": ids[0],
                    "last_id": ids[-1],
                    "has_more": has_more,
                }
            ).encode()
        )

    return respond


def _page(ids: tuple[str, ...], has_more: bool) -> _ContainerPage:
    return _ContainerPage(
        object="list",
        data=tuple(
            ContainerListResponse.model_validate(_container(container_id, container_id)) for container_id in ids
        ),
        first_id=ids[0] if ids else None,
        last_id=ids[-1] if ids else None,
        has_more=has_more,
    )


def test_container_list_returns_only_the_callers_containers_across_upstream_pages(gateway: Gateway) -> None:
    run: Final = uuid.uuid4().hex
    own_first, own_second, own_third, foreign_first, foreign_second, ownerless = (
        f"cntr_{name}_{run}" for name in ("a1", "a2", "a3", "b1", "b2", "orphan")
    )
    created_ids: Final = (own_first, own_second, own_third, foreign_first, foreign_second)
    first_page: Final = (own_first, foreign_first, own_second, ownerless)
    second_page: Final = (own_third, foreign_second)
    pages: Final = {
        None: (first_page, True),
        ownerless: (second_page, False),
        own_second: ((ownerless, *second_page), False),
    }
    provider_key: Final = f"container-key-{uuid.uuid4().hex}"
    with wire_server(_listing_upstream(created_ids, pages)) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=provider_key)
        _owner_id, owner_key = _user_key(scenario)
        _other_id, other_key = _user_key(scenario)
        owner: Final = _client(gateway, owner_key, model)
        other: Final = _client(gateway, other_key, model)
        _routable_over_alias(gateway, model)
        for index in range(2):
            _create(owner, model, index)
        for index in range(3, 5):
            _create(other, model, index)
        assert _requests(wire) == [
            ("POST", "/v1/containers", f"Bearer {provider_key}"),
            ("POST", "/v1/containers", f"Bearer {provider_key}"),
            ("POST", "/v1/containers", f"Bearer {provider_key}"),
            ("POST", "/v1/containers", f"Bearer {provider_key}"),
        ]

        first: Final = owner.containers.with_raw_response.list(limit=2, order="desc", extra_query={"model": model})
        assert _ContainerPage.model_validate_json(first.http_response.text) == _page((own_first, own_second), False), (
            first.http_response.text
        )
        assert [container.id for container in first.parse().data] == [own_first, own_second]
        auth: Final = f"Bearer {provider_key}"
        assert _requests(wire) == [
            ("GET", f"/v1/containers?limit={_OWNED_LIST_UPSTREAM_PAGE_SIZE}&order=desc", auth),
            (
                "GET",
                f"/v1/containers?after={ownerless}&limit={_OWNED_LIST_UPSTREAM_PAGE_SIZE}&order=desc",
                auth,
            ),
        ]

        _create(owner, model, 2)
        assert _requests(wire) == [("POST", "/v1/containers", f"Bearer {provider_key}")]

        listed_after_create: Final = owner.containers.with_raw_response.list(
            limit=2, order="desc", extra_query={"model": model}
        )
        assert _ContainerPage.model_validate_json(listed_after_create.http_response.text) == _page(
            (own_first, own_second), True
        ), listed_after_create.http_response.text
        assert [container.id for container in listed_after_create.parse().data] == [own_first, own_second]
        assert _requests(wire) == [
            ("GET", f"/v1/containers?limit={_OWNED_LIST_UPSTREAM_PAGE_SIZE}&order=desc", auth),
            (
                "GET",
                f"/v1/containers?after={ownerless}&limit={_OWNED_LIST_UPSTREAM_PAGE_SIZE}&order=desc",
                auth,
            ),
        ]

        alias_page: Final = gateway.client.get(
            "/containers",
            params={"limit": "2", "after": own_second, "order": "desc", "model": model},
            headers={"Authorization": f"Bearer {owner_key}"},
        )
        assert alias_page.status_code == 200, alias_page.text
        assert _ContainerPage.model_validate_json(alias_page.text) == _page((own_third,), False), alias_page.text
        assert _requests(wire) == [
            ("GET", f"/v1/containers?after={own_second}&limit={_OWNED_LIST_UPSTREAM_PAGE_SIZE}&order=desc", auth)
        ]

        admin: Final = _client(gateway, gateway.key, model).containers.with_raw_response.list(
            limit=2, order="desc", extra_query={"model": model}
        )
        assert _ContainerPage.model_validate_json(admin.http_response.text) == _page(first_page, True), (
            admin.http_response.text
        )
        assert _requests(wire) == [("GET", "/v1/containers?limit=2&order=desc", auth)]