fix(ocr-inputs): validate declared MIME, drop BMP magic sniff, normalize generic MIME

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs

@@ -290,13 +290,13 @@ fn resolve_document_mime(
         if !header.eq_ignore_ascii_case("application/octet-stream")
             && !header.eq_ignore_ascii_case("binary/octet-stream")
         {
-            return header.clone();
+            return header.to_ascii_lowercase();
         }
     }
     sniff_document_mime(bytes)
         .or_else(|| mime_from_file_name(url_path))
         .map(str::to_string)
-        .or(header_content_type)
+        .or_else(|| header_content_type.map(|header| header.to_ascii_lowercase()))
         .unwrap_or_else(|| "application/octet-stream".to_string())
 }
 
@@ -623,6 +623,26 @@ mod tests {
         );
     }
 
+    #[test]
+    fn resolve_document_mime_normalizes_specific_header_casing() {
+        assert_eq!(
+            resolve_document_mime(Some("Application/PDF".to_string()), b"%PDF-1.4", "/x"),
+            "application/pdf"
+        );
+    }
+
+    #[test]
+    fn resolve_document_mime_sniffs_when_header_is_binary_octet_stream() {
+        assert_eq!(
+            resolve_document_mime(
+                Some("Binary/Octet-Stream".to_string()),
+                b"%PDF-1.4 payload",
+                "/x"
+            ),
+            "application/pdf"
+        );
+    }
+
     #[test]
     fn resolve_document_mime_sniffs_when_header_is_octet_stream() {
         assert_eq!(

litellm-rust/crates/core/src/ocr/mime.rs

@@ -18,9 +18,6 @@ pub fn sniff_document_mime(bytes: &[u8]) -> Option<&'static str> {
     {
         return Some("image/tiff");
     }
-    if bytes.starts_with(b"BM") {
-        return Some("image/bmp");
-    }
     None
 }
 
@@ -71,7 +68,12 @@ mod tests {
             sniff_document_mime(&[0x4d, 0x4d, 0x00, 0x2a]),
             Some("image/tiff")
         );
-        assert_eq!(sniff_document_mime(b"BMxxxx"), Some("image/bmp"));
+    }
+
+    #[test]
+    fn sniff_document_mime_does_not_detect_bmp_by_magic_bytes() {
+        assert_eq!(sniff_document_mime(b"BMxxxx"), None);
+        assert_eq!(sniff_document_mime(b"BM\x36\x00\x00\x00random"), None);
     }
 
     #[test]

litellm/ocr/main.py

@@ -372,6 +372,8 @@ async def aocr(
 
 _MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$")
 
+_GENERIC_MIME_TYPES = frozenset({"application/octet-stream", "binary/octet-stream"})
+
 _RUST_BRIDGE_INTERNAL_PARAMS = {"original_generic_function"}
 
 _MIME_TYPE_MAP = {
@@ -414,11 +416,27 @@ def _sniff_mime_type_from_bytes(data: bytes) -> str | None:
         return "image/webp"
     if data.startswith((b"II*\x00", b"MM\x00*")):
         return "image/tiff"
-    if data.startswith(b"BM"):
-        return "image/bmp"
     return None
 
 
+def _normalize_mime_type(value: str) -> str | None:
+    normalized = value.split(";", 1)[0].strip().lower()
+    if not _MIME_PATTERN.match(normalized):
+        return None
+    return normalized
+
+
+def _resolve_declared_mime_type(raw: object) -> str | None:
+    if raw is None:
+        return None
+    if not isinstance(raw, str):
+        raise ValueError("OCR document 'mime_type' must be a string when provided")
+    normalized = _normalize_mime_type(raw)
+    if normalized is None:
+        raise ValueError("OCR document 'mime_type' is not a valid MIME type")
+    return normalized
+
+
 def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]:
     """
     Convert a file-type document dict to a document_url-type document dict
@@ -488,25 +506,26 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str,
     if not file_bytes:
         raise ValueError("File is empty or could not be read")
 
-    raw_mime = document.get("mime_type")
-    explicit_mime = raw_mime if isinstance(raw_mime, str) else None
-    resolved_mime = next(
+    declared_mime = (
+        _resolve_declared_mime_type(cast(object, document["mime_type"]))
+        if "mime_type" in document
+        else None
+    )
+    filename_mime = _normalize_mime_type(mime_type)
+    specific_mime = next(
         (
             candidate
-            for candidate in (explicit_mime, mime_type)
-            if candidate and candidate != "application/octet-stream"
+            for candidate in (declared_mime, filename_mime)
+            if candidate is not None and candidate not in _GENERIC_MIME_TYPES
         ),
         None,
     )
     mime_type = (
-        resolved_mime
+        specific_mime
         or _sniff_mime_type_from_bytes(file_bytes)
         or "application/octet-stream"
     )
 
-    if not _MIME_PATTERN.match(mime_type):
-        raise ValueError(f"Invalid MIME type: {mime_type}")
-
     base64_data = base64.b64encode(file_bytes).decode("utf-8")
     data_uri = f"data:{mime_type};base64,{base64_data}"
 

litellm/proxy/ocr_endpoints/endpoints.py

@@ -27,10 +27,10 @@ def _build_document_from_upload(
     Delegates to convert_file_document_to_url_document after resolving MIME type
     from the upload's content_type header or filename.
     """
-    mime_type = content_type.split(";")[0].strip() if content_type else None
-    if not mime_type or mime_type == "application/octet-stream":
-        if filename:
-            mime_type = get_mime_type(filename)
+    mime_type = content_type.split(";")[0].strip().lower() if content_type else None
+    generic_mime_types = ("application/octet-stream", "binary/octet-stream")
+    if (not mime_type or mime_type in generic_mime_types) and filename:
+        mime_type = get_mime_type(filename)
 
     return convert_file_document_to_url_document(
         {

tests/e2e/gateway/test_ocr_rust_e2e.py

@@ -11,13 +11,12 @@ from __future__ import annotations
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
 
 import httpx
 import pytest
 import yaml
 
-from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
 
 TEST_PDF_URL = (
     "https://cdn.jsdelivr.net/gh/BerriAI/litellm"
@@ -34,12 +33,12 @@ REPO_ROOT = Path(__file__).resolve().parents[3]
 FIXTURE_PDF = REPO_ROOT / "tests" / "llm_translation" / "fixtures" / "dummy.pdf"
 FIXTURE_IMAGE = REPO_ROOT / "tests" / "image_gen_tests" / "test_image.png"
 
-RUST_OCR_UPLOAD_CASES = [
-    pytest.param(FIXTURE_PDF, "application/pdf", "document_url", id="pdf_octet_stream"),
-    pytest.param(FIXTURE_IMAGE, "image/png", "image_url", id="image_octet_stream"),
-]
+RUST_OCR_UPLOAD_CASES = (
+    pytest.param(FIXTURE_PDF, id="pdf_octet_stream"),
+    pytest.param(FIXTURE_IMAGE, id="image_octet_stream"),
+)
 
-RUST_OCR_GATEWAY_CASES = [
+RUST_OCR_GATEWAY_CASES = (
     pytest.param(
         "rust-ocr-mistral",
         {"type": "document_url", "document_url": TEST_PDF_URL},
@@ -68,7 +67,7 @@ RUST_OCR_GATEWAY_CASES = [
         },
         id="vertex_deepseek",
     ),
-]
+)
 
 CONFIG_PATH = Path(__file__).with_name("litellm-config.yml")
 
@@ -103,7 +102,9 @@ class OcrGateway:
                 json={"model": model, "document": document},
             )
 
-    def ocr_upload(self, model: str, content: bytes, upload_name: str) -> httpx.Response:
+    def ocr_upload(
+        self, model: str, content: bytes, upload_name: str
+    ) -> httpx.Response:
         with httpx.Client(
             timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
         ) as client:
@@ -135,13 +136,13 @@ def resources() -> OcrResources:
     )
 
 
-def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None:
-    assert response_json["object"] == "ocr"
-    assert response_json["model"]
-    assert isinstance(response_json["pages"], list)
-    assert len(response_json["pages"]) > 0
-    assert "index" in response_json["pages"][0]
-    assert "markdown" in response_json["pages"][0]
+def _assert_ocr_response(response: httpx.Response) -> None:
+    assert response.status_code == 200, response.text
+    parsed = OCRResponse.model_validate(response.json())
+    assert parsed.object == "ocr"
+    assert parsed.model
+    assert len(parsed.pages) > 0
+    assert any(page.markdown.strip() for page in parsed.pages)
 
 
 class TestRustOcrGateway:
@@ -166,29 +167,17 @@ class TestRustOcrGateway:
     ) -> None:
         response = resources.gateway.ocr(model, document)
 
-        assert response.status_code == 200, response.text
-        _assert_ocr_response_shape(response.json())
+        _assert_ocr_response(response)
 
-    @pytest.mark.parametrize(
-        ("fixture_path", "expected_mime", "expected_document_type"),
-        RUST_OCR_UPLOAD_CASES,
-    )
+    @pytest.mark.parametrize("fixture_path", RUST_OCR_UPLOAD_CASES)
     def test_rust_ocr_octet_stream_upload_response(
         self,
         resources: OcrResources,
         fixture_path: Path,
-        expected_mime: str,
-        expected_document_type: str,
     ) -> None:
-        if not fixture_path.is_file():
-            pytest.skip(f"Missing OCR fixture: {fixture_path}")
+        assert fixture_path.is_file(), f"Missing committed OCR fixture: {fixture_path}"
 
         content = fixture_path.read_bytes()
-        document = _build_document_from_upload(content, "document", "application/octet-stream")
-        assert document["type"] == expected_document_type
-        assert document[expected_document_type].startswith(f"data:{expected_mime};base64,")
-
         response = resources.gateway.ocr_upload("rust-ocr-mistral", content, "document")
 
-        assert response.status_code == 200, response.text
-        _assert_ocr_response_shape(response.json())
+        _assert_ocr_response(response)

tests/test_litellm/ocr/test_ocr_file_input.py

@@ -85,8 +85,9 @@ class TestSniffMimeTypeFromBytes:
         assert _sniff_mime_type_from_bytes(b"II*\x00 rest") == "image/tiff"
         assert _sniff_mime_type_from_bytes(b"MM\x00* rest") == "image/tiff"
 
-    def test_should_detect_bmp_from_magic_bytes(self):
-        assert _sniff_mime_type_from_bytes(b"BM rest") == "image/bmp"
+    def test_should_not_sniff_bmp_from_magic_bytes(self) -> None:
+        assert _sniff_mime_type_from_bytes(b"BM rest") is None
+        assert _sniff_mime_type_from_bytes(b"BM\x00\x00\x00\x00random") is None
 
     def test_should_return_none_for_unknown_prefix(self):
         assert _sniff_mime_type_from_bytes(b"not a known file") is None
@@ -333,18 +334,119 @@ class TestConvertFileDocumentToUrlDocument:
         with pytest.raises(ValueError, match="Unsupported file input type"):
             convert_file_document_to_url_document({"type": "file", "file": 12345})
 
-    def test_should_raise_error_for_invalid_mime_type(self):
-        """MIME types with special characters should be rejected."""
+    def test_should_reject_malformed_mime_type(self) -> None:
         content = b"some content"
-        with pytest.raises(ValueError, match="Invalid MIME type"):
+        with pytest.raises(ValueError, match="not a valid MIME type"):
             convert_file_document_to_url_document(
                 {
                     "type": "file",
                     "file": content,
-                    "mime_type": "text/html; charset=utf-8\nX-Injected: true",
+                    "mime_type": "text/html\r\nX-Injected: true",
                 }
             )
 
+    def test_should_reject_blank_mime_type(self) -> None:
+        content = b"%PDF-1.4 content"
+        with pytest.raises(ValueError, match="not a valid MIME type"):
+            convert_file_document_to_url_document(
+                {"type": "file", "file": content, "mime_type": "   "}
+            )
+
+    def test_should_reject_non_string_mime_type(self) -> None:
+        content = b"%PDF-1.4 content"
+        with pytest.raises(ValueError, match="must be a string"):
+            convert_file_document_to_url_document(
+                {"type": "file", "file": content, "mime_type": 123}
+            )
+
+    def test_should_not_echo_raw_mime_value_in_error(self) -> None:
+        content = b"some content"
+        secret_marker = "X-Injected: super-secret"
+        with pytest.raises(ValueError) as exc_info:
+            convert_file_document_to_url_document(
+                {
+                    "type": "file",
+                    "file": content,
+                    "mime_type": f"text/html\r\n{secret_marker}",
+                }
+            )
+        assert secret_marker not in str(exc_info.value)
+
+    def test_should_strip_parameters_from_explicit_mime_type(self) -> None:
+        content = b"raw bytes"
+        result = convert_file_document_to_url_document(
+            {
+                "type": "file",
+                "file": content,
+                "mime_type": "application/pdf; charset=utf-8",
+            }
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_normalize_explicit_mime_casing(self) -> None:
+        content = b"raw bytes"
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content, "mime_type": "Application/PDF"}
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_sniff_when_explicit_mime_is_uppercase_octet_stream(self) -> None:
+        content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n"
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content, "mime_type": "Application/Octet-Stream"}
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_sniff_when_explicit_mime_is_binary_octet_stream(self) -> None:
+        content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n"
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content, "mime_type": "Binary/Octet-Stream"}
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_treat_generic_octet_stream_with_parameters_as_ambiguous(
+        self,
+    ) -> None:
+        content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n"
+        result = convert_file_document_to_url_document(
+            {
+                "type": "file",
+                "file": content,
+                "mime_type": "application/octet-stream; charset=binary",
+            }
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_not_sniff_bmp_from_raw_bytes(self) -> None:
+        content = b"BM\x36\x00\x00\x00 fake bitmap header"
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content}
+        )
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith(
+            "data:application/octet-stream;base64,"
+        )
+
+    def test_should_resolve_bmp_from_extension(self) -> None:
+        content = b"BM\x36\x00\x00\x00 fake bitmap header"
+        with tempfile.NamedTemporaryFile(suffix=".bmp", delete=False) as f:
+            f.write(content)
+            f.flush()
+            tmp_path = Path(f.name)
+        try:
+            result = convert_file_document_to_url_document(
+                {"type": "file", "file": tmp_path}
+            )
+            assert result["type"] == "image_url"
+            assert result["image_url"].startswith("data:image/bmp;base64,")
+        finally:
+            os.unlink(str(tmp_path))
+
     def test_should_override_mime_type_for_pathlib_path(self):
         """Explicit mime_type should override auto-detection from extension."""
         content = b"some content"
@@ -493,6 +595,30 @@ class TestBuildDocumentFromUpload:
         assert result["type"] == "image_url"
         assert result["image_url"].startswith("data:image/png;base64,")
 
+    def test_should_detect_mime_from_filename_for_uppercase_octet_stream(self) -> None:
+        content = b"pdf content"
+
+        result = self._build(
+            file_content=content,
+            filename="report.pdf",
+            content_type="APPLICATION/OCTET-STREAM",
+        )
+
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_detect_mime_from_filename_for_binary_octet_stream(self) -> None:
+        content = b"png content"
+
+        result = self._build(
+            file_content=content,
+            filename="image.png",
+            content_type="binary/octet-stream",
+        )
+
+        assert result["type"] == "image_url"
+        assert result["image_url"].startswith("data:image/png;base64,")
+
 
 class TestProxySecurityGuard:
     """Test that the proxy rejects type='file' documents in JSON requests