From f598cfe3a67fcbce40755c3e7eaa754e60be2cbd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:02:39 +0000 Subject: [PATCH] feat(rust-gateway): sniff upload MIME by magic bytes and preserve reducto ids Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../crates/ai-gateway/src/constants.rs | 25 ++-- .../crates/ai-gateway/src/routes/ocr/mod.rs | 118 +++++++++++++++++- .../ai-gateway/src/routes/ocr/transport.rs | 60 ++++++--- litellm-rust/crates/core/src/constants.rs | 8 ++ litellm-rust/crates/core/src/lib.rs | 1 + litellm-rust/crates/core/src/ocr/mime.rs | 93 ++++++++++++++ litellm-rust/crates/core/src/ocr/mod.rs | 1 + 7 files changed, 277 insertions(+), 29 deletions(-) create mode 100644 litellm-rust/crates/core/src/constants.rs create mode 100644 litellm-rust/crates/core/src/ocr/mime.rs diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 8a18dd14260..0abc7b04592 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -5,6 +5,11 @@ //! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env //! read + fallback happens at the host/config layer. +use litellm_core::constants::{ + MIME_APPLICATION_OCTET_STREAM, MIME_APPLICATION_PDF, MIME_IMAGE_BMP, MIME_IMAGE_GIF, + MIME_IMAGE_JPEG, MIME_IMAGE_PNG, MIME_IMAGE_TIFF, MIME_IMAGE_WEBP, +}; + /// Default LiteLLM control-plane base URL for request-log egress when /// `LITELLM_PROXY_BASE_URL` is unset. pub(crate) const DEFAULT_PROXY_BASE_URL: &str = "http://localhost:4000"; @@ -28,18 +33,18 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; /// Provider attributed to realtime sessions in the logging payload. pub(crate) const DEFAULT_PROVIDER: &str = "openai"; -pub(crate) const DEFAULT_UPLOAD_MIME_TYPE: &str = "application/octet-stream"; +pub(crate) const DEFAULT_UPLOAD_MIME_TYPE: &str = MIME_APPLICATION_OCTET_STREAM; pub(crate) const MAX_OCR_REQUEST_BYTES: usize = 100 * 1024 * 1024; pub(crate) const OCR_UPLOAD_MIME_BY_EXTENSION: &[(&str, &str)] = &[ - ("pdf", "application/pdf"), - ("png", "image/png"), - ("jpg", "image/jpeg"), - ("jpeg", "image/jpeg"), - ("gif", "image/gif"), - ("webp", "image/webp"), - ("tiff", "image/tiff"), - ("tif", "image/tiff"), - ("bmp", "image/bmp"), + ("pdf", MIME_APPLICATION_PDF), + ("png", MIME_IMAGE_PNG), + ("jpg", MIME_IMAGE_JPEG), + ("jpeg", MIME_IMAGE_JPEG), + ("gif", MIME_IMAGE_GIF), + ("webp", MIME_IMAGE_WEBP), + ("tiff", MIME_IMAGE_TIFF), + ("tif", MIME_IMAGE_TIFF), + ("bmp", MIME_IMAGE_BMP), ]; diff --git a/litellm-rust/crates/ai-gateway/src/routes/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/ocr/mod.rs index 05b3553fda1..a3de414dd9b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/ocr/mod.rs @@ -129,7 +129,7 @@ fn error_response(error: &CoreError) -> Response { CoreError::Auth(_) => ( StatusCode::UNAUTHORIZED, "authentication_error", - error.to_string(), + "authentication failed".to_string(), ), CoreError::Routing(_) => (StatusCode::NOT_FOUND, "not_found_error", error.to_string()), CoreError::Http { status, .. } => { @@ -197,6 +197,61 @@ mod tests { String::from_utf8_lossy(&request).into_owned() } + async fn read_full_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 4096]; + let mut body_start: Option = None; + let mut content_length = 0_usize; + loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + if body_start.is_none() { + if let Some(pos) = request.windows(4).position(|window| window == b"\r\n\r\n") { + let start = pos + 4; + body_start = Some(start); + let headers = String::from_utf8_lossy(&request[..pos]).to_ascii_lowercase(); + content_length = headers + .lines() + .find_map(|line| line.strip_prefix("content-length:")) + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(0); + } + } + if let Some(start) = body_start { + if request.len() >= start + content_length { + break; + } + } + } + String::from_utf8_lossy(&request).into_owned() + } + + async fn spawn_mock_upstream_capture() -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("binds upstream"); + let addr = listener.local_addr().expect("upstream addr"); + let handle = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_full_http_request(&mut socket).await; + let body = r#"{"pages":[{"index":0,"markdown":"hello ocr"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + (format!("http://{addr}"), handle) + } + async fn spawn_mock_upstream() -> (String, tokio::task::JoinHandle) { let listener = TcpListener::bind("127.0.0.1:0") .await @@ -374,6 +429,67 @@ mod tests { assert!(upstream_request.starts_with("POST"), "{upstream_request}"); } + #[tokio::test] + async fn multipart_unnamed_octet_stream_pdf_is_sniffed() { + let (upstream, upstream_handle) = spawn_mock_upstream_capture().await; + let addr = serve(app_with_deployment(&upstream)).await; + + let form = reqwest::multipart::Form::new() + .text("model", "rust-ocr-mistral") + .part( + "file", + reqwest::multipart::Part::bytes(b"%PDF-1.7 minimal pdf bytes".to_vec()) + .mime_str("application/octet-stream") + .expect("mime"), + ); + + let response = reqwest::Client::new() + .post(format!("http://{addr}/v1/ocr")) + .bearer_auth(MASTER_KEY) + .multipart(form) + .send() + .await + .expect("request sent"); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + let upstream_request = upstream_handle.await.expect("upstream served"); + assert!( + upstream_request.contains("data:application/pdf;base64,"), + "unnamed octet-stream PDF must be sniffed to application/pdf: {upstream_request}" + ); + } + + #[tokio::test] + async fn multipart_unnamed_octet_stream_image_is_sniffed() { + let (upstream, upstream_handle) = spawn_mock_upstream_capture().await; + let addr = serve(app_with_deployment(&upstream)).await; + + let png = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x01]; + let form = reqwest::multipart::Form::new() + .text("model", "rust-ocr-mistral") + .part( + "file", + reqwest::multipart::Part::bytes(png) + .mime_str("application/octet-stream") + .expect("mime"), + ); + + let response = reqwest::Client::new() + .post(format!("http://{addr}/v1/ocr")) + .bearer_auth(MASTER_KEY) + .multipart(form) + .send() + .await + .expect("request sent"); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + let upstream_request = upstream_handle.await.expect("upstream served"); + assert!( + upstream_request.contains("data:image/png;base64,"), + "unnamed octet-stream PNG must be sniffed to image/png: {upstream_request}" + ); + } + #[tokio::test] async fn upstream_error_status_is_propagated_without_leaking_provider_body() { let upstream = spawn_mock_upstream_error().await; diff --git a/litellm-rust/crates/ai-gateway/src/routes/ocr/transport.rs b/litellm-rust/crates/ai-gateway/src/routes/ocr/transport.rs index dcdf73a5a69..c106c0cf5b6 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/ocr/transport.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/ocr/transport.rs @@ -3,6 +3,7 @@ use std::time::Duration; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::Engine; use litellm_core::error::CoreError; +use litellm_core::ocr::mime::sniff_mime; use litellm_core::CoreResult; use serde::Deserialize; use serde_json::{json, Map, Value}; @@ -31,14 +32,6 @@ impl OcrDocument { )) } }; - if url.starts_with("reducto://") { - return Err(CoreError::InvalidRequest( - "reducto:// file IDs are not accepted through the OCR API; upload the file in \ - the same request via multipart/form-data with a 'file' field, or pass an \ - inline base64 data URI as the document URL" - .to_string(), - )); - } Ok(json!({ "type": field, field: url })) } } @@ -95,7 +88,7 @@ fn mime_from_filename(filename: &str) -> Option<&'static str> { .map(|(_, mime)| *mime) } -fn resolve_upload_mime(content_type: Option<&str>, filename: Option<&str>) -> String { +fn resolve_upload_mime(bytes: &[u8], content_type: Option<&str>, filename: Option<&str>) -> String { let declared = content_type .and_then(|value| value.split(';').next()) .map(str::trim) @@ -103,6 +96,9 @@ fn resolve_upload_mime(content_type: Option<&str>, filename: Option<&str>) -> St if let Some(declared) = declared { return declared.to_string(); } + if let Some(sniffed) = sniff_mime(bytes) { + return sniffed.to_string(); + } filename .and_then(mime_from_filename) .unwrap_or(DEFAULT_UPLOAD_MIME_TYPE) @@ -119,7 +115,7 @@ pub fn build_upload_document( "uploaded file is empty".to_string(), )); } - let mime = resolve_upload_mime(content_type, filename); + let mime = resolve_upload_mime(&bytes, content_type, filename); let data_uri = format!("data:{mime};base64,{}", BASE64_STANDARD.encode(&bytes)); let field = if mime.starts_with("image/") { "image_url" @@ -221,15 +217,13 @@ mod tests { } #[test] - fn rejects_reducto_file_id_over_json() { - let err = parse_json_body( - br#"{"model":"m","document":{"type":"document_url","document_url":"reducto://abc"}}"#, + fn preserves_reducto_file_id_over_json() { + let call = parse_json_body( + br#"{"model":"m","document":{"type":"document_url","document_url":"reducto://abc123"}}"#, ) - .expect_err("reducto id rejected"); - match err { - CoreError::InvalidRequest(message) => assert!(message.contains("reducto://")), - other => panic!("expected InvalidRequest, got {other:?}"), - } + .expect("reducto id preserved"); + assert_eq!(call.document["type"], "document_url"); + assert_eq!(call.document["document_url"], "reducto://abc123"); } #[test] @@ -255,6 +249,36 @@ mod tests { .starts_with("data:application/pdf;base64,")); } + #[test] + fn upload_sniffs_pdf_from_bytes_when_unnamed_octet_stream() { + let document = build_upload_document( + b"%PDF-1.7 minimal".to_vec(), + None, + Some("application/octet-stream"), + ) + .expect("builds document"); + assert_eq!(document["type"], "document_url"); + assert!(document["document_url"] + .as_str() + .expect("data uri") + .starts_with("data:application/pdf;base64,")); + } + + #[test] + fn upload_sniffs_png_from_bytes_when_unnamed_and_no_content_type() { + let document = build_upload_document( + vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00], + None, + None, + ) + .expect("builds document"); + assert_eq!(document["type"], "image_url"); + assert!(document["image_url"] + .as_str() + .expect("data uri") + .starts_with("data:image/png;base64,")); + } + #[test] fn upload_infers_mime_from_filename_when_octet_stream() { let document = build_upload_document( diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs new file mode 100644 index 00000000000..4aff46a7a7e --- /dev/null +++ b/litellm-rust/crates/core/src/constants.rs @@ -0,0 +1,8 @@ +pub const MIME_APPLICATION_PDF: &str = "application/pdf"; +pub const MIME_APPLICATION_OCTET_STREAM: &str = "application/octet-stream"; +pub const MIME_IMAGE_PNG: &str = "image/png"; +pub const MIME_IMAGE_JPEG: &str = "image/jpeg"; +pub const MIME_IMAGE_GIF: &str = "image/gif"; +pub const MIME_IMAGE_WEBP: &str = "image/webp"; +pub const MIME_IMAGE_TIFF: &str = "image/tiff"; +pub const MIME_IMAGE_BMP: &str = "image/bmp"; diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 2ac479cc725..e0e15d8071c 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod constants; pub mod error; pub mod ocr; pub mod providers; diff --git a/litellm-rust/crates/core/src/ocr/mime.rs b/litellm-rust/crates/core/src/ocr/mime.rs new file mode 100644 index 00000000000..79e6deb95bf --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/mime.rs @@ -0,0 +1,93 @@ +use crate::constants::{ + MIME_APPLICATION_PDF, MIME_IMAGE_BMP, MIME_IMAGE_GIF, MIME_IMAGE_JPEG, MIME_IMAGE_PNG, + MIME_IMAGE_TIFF, MIME_IMAGE_WEBP, +}; + +pub fn sniff_mime(bytes: &[u8]) -> Option<&'static str> { + if bytes.starts_with(b"%PDF") { + return Some(MIME_APPLICATION_PDF); + } + if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) { + return Some(MIME_IMAGE_PNG); + } + if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) { + return Some(MIME_IMAGE_JPEG); + } + if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") { + return Some(MIME_IMAGE_GIF); + } + if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" { + return Some(MIME_IMAGE_WEBP); + } + if bytes.starts_with(&[0x49, 0x49, 0x2A, 0x00]) || bytes.starts_with(&[0x4D, 0x4D, 0x00, 0x2A]) + { + return Some(MIME_IMAGE_TIFF); + } + if bytes.starts_with(b"BM") { + return Some(MIME_IMAGE_BMP); + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sniffs_pdf() { + assert_eq!(sniff_mime(b"%PDF-1.7\n..."), Some(MIME_APPLICATION_PDF)); + } + + #[test] + fn sniffs_png() { + assert_eq!( + sniff_mime(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00]), + Some(MIME_IMAGE_PNG) + ); + } + + #[test] + fn sniffs_jpeg() { + assert_eq!(sniff_mime(&[0xFF, 0xD8, 0xFF, 0xE0]), Some(MIME_IMAGE_JPEG)); + } + + #[test] + fn sniffs_gif() { + assert_eq!(sniff_mime(b"GIF87a...."), Some(MIME_IMAGE_GIF)); + assert_eq!(sniff_mime(b"GIF89a...."), Some(MIME_IMAGE_GIF)); + } + + #[test] + fn sniffs_webp() { + let mut bytes = b"RIFF".to_vec(); + bytes.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); + bytes.extend_from_slice(b"WEBP"); + assert_eq!(sniff_mime(&bytes), Some(MIME_IMAGE_WEBP)); + } + + #[test] + fn does_not_sniff_riff_without_webp() { + let mut bytes = b"RIFF".to_vec(); + bytes.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); + bytes.extend_from_slice(b"WAVE"); + assert_eq!(sniff_mime(&bytes), None); + } + + #[test] + fn sniffs_tiff_both_byte_orders() { + assert_eq!(sniff_mime(&[0x49, 0x49, 0x2A, 0x00]), Some(MIME_IMAGE_TIFF)); + assert_eq!(sniff_mime(&[0x4D, 0x4D, 0x00, 0x2A]), Some(MIME_IMAGE_TIFF)); + } + + #[test] + fn sniffs_bmp() { + assert_eq!(sniff_mime(b"BM...."), Some(MIME_IMAGE_BMP)); + } + + #[test] + fn returns_none_for_unknown_and_truncated() { + assert_eq!(sniff_mime(&[0x00, 0x01, 0x02, 0x03]), None); + assert_eq!(sniff_mime(b""), None); + assert_eq!(sniff_mime(&[0x89, b'P', b'N', b'G']), None); + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index ec2fbb969a6..c0a0d9bb565 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,3 @@ +pub mod mime; pub mod transformation; pub mod types;