mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 02:47:45 +08:00
feat(security): harden gateway boundaries and usage policies
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change. Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
@@ -1,3 +1,96 @@
|
||||
use url::Url;
|
||||
|
||||
pub mod request;
|
||||
pub mod spec;
|
||||
pub mod stream;
|
||||
|
||||
pub(crate) const MAX_OPENAI_IMAGE_DATA_BYTES: usize = 64 * 1024 * 1024;
|
||||
pub(crate) const MAX_OPENAI_IMAGE_EXTERNAL_URL_BYTES: usize = 64 * 1024;
|
||||
pub(crate) const MAX_OPENAI_IMAGE_REVISED_PROMPT_BYTES: usize = 256 * 1024;
|
||||
|
||||
pub(crate) fn is_safe_openai_image_base64_payload(value: &str) -> bool {
|
||||
if value.is_empty()
|
||||
|| value.len() > MAX_OPENAI_IMAGE_DATA_BYTES
|
||||
|| value.len() % 4 == 1
|
||||
|| value
|
||||
.bytes()
|
||||
.any(|byte| byte.is_ascii_whitespace() || byte.is_ascii_control())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let bytes = value.as_bytes();
|
||||
let first_padding = bytes.iter().position(|byte| *byte == b'=');
|
||||
if let Some(index) = first_padding {
|
||||
let padding = bytes.len() - index;
|
||||
if padding > 2 || bytes[index..].iter().any(|byte| *byte != b'=') || bytes.len() % 4 != 0 {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
bytes[..first_padding.unwrap_or(bytes.len())]
|
||||
.iter()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(*byte, b'+' | b'/'))
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_openai_image_output_format(value: &str) -> Option<&'static str> {
|
||||
let value = value.trim();
|
||||
if value.eq_ignore_ascii_case("png") {
|
||||
Some("png")
|
||||
} else if value.eq_ignore_ascii_case("jpeg") || value.eq_ignore_ascii_case("jpg") {
|
||||
Some("jpeg")
|
||||
} else if value.eq_ignore_ascii_case("webp") {
|
||||
Some("webp")
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn bounded_openai_image_revised_prompt(value: &str) -> Option<&str> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty() && value.len() <= MAX_OPENAI_IMAGE_REVISED_PROMPT_BYTES).then_some(value)
|
||||
}
|
||||
|
||||
pub(crate) fn parse_safe_openai_image_data_url(value: &str) -> Option<(&'static str, &str)> {
|
||||
let (metadata, payload) = value.trim().split_once(',')?;
|
||||
let mime_type = metadata.strip_prefix("data:")?.strip_suffix(";base64")?;
|
||||
let mime_type = safe_openai_image_mime_type(mime_type.trim())?;
|
||||
(!payload.is_empty() && is_safe_openai_image_base64_payload(payload))
|
||||
.then_some((mime_type, payload))
|
||||
}
|
||||
|
||||
pub(crate) fn safe_openai_image_mime_type(value: &str) -> Option<&'static str> {
|
||||
if value.eq_ignore_ascii_case("image/png") {
|
||||
Some("image/png")
|
||||
} else if value.eq_ignore_ascii_case("image/jpeg") || value.eq_ignore_ascii_case("image/jpg") {
|
||||
Some("image/jpeg")
|
||||
} else if value.eq_ignore_ascii_case("image/webp") {
|
||||
Some("image/webp")
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn sanitize_openai_image_source_url(value: &str) -> Option<String> {
|
||||
let value = value.trim();
|
||||
if value.is_empty()
|
||||
|| value
|
||||
.chars()
|
||||
.any(|character| character.is_ascii_control() || character.is_whitespace())
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if let Some((mime_type, payload)) = parse_safe_openai_image_data_url(value) {
|
||||
return Some(format!("data:{mime_type};base64,{payload}"));
|
||||
}
|
||||
if value.len() > MAX_OPENAI_IMAGE_EXTERNAL_URL_BYTES {
|
||||
return None;
|
||||
}
|
||||
let parsed = Url::parse(value).ok()?;
|
||||
if !matches!(parsed.scheme(), "http" | "https")
|
||||
|| parsed.host_str().is_none()
|
||||
|| !parsed.username().is_empty()
|
||||
|| parsed.password().is_some()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(value.to_string())
|
||||
}
|
||||
|
||||
@@ -71,6 +71,9 @@ impl OpenAiImageNormalizeOptions {
|
||||
|
||||
pub const CHATGPT_WEB_IMAGE_MAX_AREA: u64 = 1_500_000;
|
||||
pub const OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 10;
|
||||
const MAX_MULTIPART_PARTS: usize = 128;
|
||||
const MAX_MULTIPART_PART_HEADER_BYTES: usize = 64 * 1024;
|
||||
const MAX_MULTIPART_BODY_BYTES: usize = 256 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ChatGptWebImageRequestError {
|
||||
@@ -1653,34 +1656,69 @@ fn parse_multipart_fields_from_base64(
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())?;
|
||||
let boundary = multipart_boundary(content_type)?;
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
let body_bytes =
|
||||
decode_multipart_body_base64_with_limit(body_base64, MAX_MULTIPART_BODY_BYTES)?;
|
||||
Some(parse_multipart_fields(&body_bytes, boundary.as_str()))
|
||||
}
|
||||
|
||||
fn decode_multipart_body_base64_with_limit(
|
||||
body_base64: &str,
|
||||
decoded_limit: usize,
|
||||
) -> Option<Vec<u8>> {
|
||||
let max_encoded_len = decoded_limit
|
||||
.checked_add(2)
|
||||
.and_then(|value| value.checked_div(3))
|
||||
.and_then(|value| value.checked_mul(4))
|
||||
.unwrap_or(usize::MAX);
|
||||
if body_base64.len() > max_encoded_len {
|
||||
return None;
|
||||
}
|
||||
let body = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.ok()?;
|
||||
Some(parse_multipart_fields(&body_bytes, boundary.as_str()))
|
||||
(body.len() <= decoded_limit).then_some(body)
|
||||
}
|
||||
|
||||
fn parse_multipart_fields(body: &[u8], boundary: &str) -> Vec<MultipartField> {
|
||||
let delimiter = format!("--{boundary}").into_bytes();
|
||||
let mut parts = Vec::new();
|
||||
let mut cursor = 0usize;
|
||||
let mut part_count = 0usize;
|
||||
|
||||
while let Some(index) = find_subslice(&body[cursor..], &delimiter) {
|
||||
while let Some(index) = find_multipart_boundary(&body[cursor..], &delimiter, true) {
|
||||
let start = cursor + index + delimiter.len();
|
||||
if body.get(start..start + 2) == Some(b"--") {
|
||||
let closing_suffix = body.get(start + 2..).unwrap_or_default();
|
||||
if !(closing_suffix.is_empty() || closing_suffix.starts_with(b"\r\n")) {
|
||||
return Vec::new();
|
||||
}
|
||||
break;
|
||||
}
|
||||
part_count = part_count.saturating_add(1);
|
||||
if part_count > MAX_MULTIPART_PARTS {
|
||||
return Vec::new();
|
||||
}
|
||||
let mut part = &body[start..];
|
||||
if part.starts_with(b"\r\n") {
|
||||
part = &part[2..];
|
||||
}
|
||||
let Some(next) = find_subslice(part, &delimiter) else {
|
||||
break;
|
||||
// Do not return fields parsed before a truncated part. Callers use
|
||||
// an empty result as the invalid-multipart signal, so retaining a
|
||||
// prefix would turn malformed input into an accepted request.
|
||||
let Some(next) = find_multipart_boundary(part, &delimiter, false) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let raw = &part[..next];
|
||||
let raw = raw.strip_suffix(b"\r\n").unwrap_or(raw);
|
||||
if let Some(field) = parse_multipart_field(raw) {
|
||||
parts.push(field);
|
||||
if find_subslice(raw, b"\r\n\r\n")
|
||||
.is_some_and(|header_end| header_end > MAX_MULTIPART_PART_HEADER_BYTES)
|
||||
{
|
||||
return Vec::new();
|
||||
}
|
||||
let Some(field) = parse_multipart_field(raw) else {
|
||||
return Vec::new();
|
||||
};
|
||||
parts.push(field);
|
||||
cursor = start + next;
|
||||
}
|
||||
|
||||
@@ -1688,34 +1726,219 @@ fn parse_multipart_fields(body: &[u8], boundary: &str) -> Vec<MultipartField> {
|
||||
}
|
||||
|
||||
fn multipart_boundary(content_type: &str) -> Option<String> {
|
||||
content_type.split(';').find_map(|segment| {
|
||||
let (key, value) = segment.trim().split_once('=')?;
|
||||
if !key.trim().eq_ignore_ascii_case("boundary") {
|
||||
let segments = split_multipart_content_type_parameters(content_type)?;
|
||||
let media_type = segments.first()?.trim();
|
||||
if !media_type.eq_ignore_ascii_case("multipart/form-data") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut boundary = None;
|
||||
let mut seen_keys = Vec::new();
|
||||
for segment in segments.into_iter().skip(1) {
|
||||
let segment = segment.trim();
|
||||
if segment.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let boundary = value.trim().trim_matches('"').trim();
|
||||
(!boundary.is_empty()).then(|| boundary.to_string())
|
||||
})
|
||||
let (raw_key, raw_value) = segment.split_once('=')?;
|
||||
let key = raw_key.trim();
|
||||
if key.is_empty() || !key.as_bytes().iter().copied().all(is_http_token_byte) {
|
||||
return None;
|
||||
}
|
||||
if seen_keys
|
||||
.iter()
|
||||
.any(|seen: &String| seen.eq_ignore_ascii_case(key))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
seen_keys.push(key.to_ascii_lowercase());
|
||||
|
||||
let (value, had_escape) = parse_multipart_content_type_parameter_value(raw_value.trim())?;
|
||||
if !key.eq_ignore_ascii_case("boundary") {
|
||||
continue;
|
||||
}
|
||||
if had_escape || !is_valid_multipart_boundary(&value) {
|
||||
return None;
|
||||
}
|
||||
boundary = Some(value);
|
||||
}
|
||||
|
||||
boundary
|
||||
}
|
||||
|
||||
fn split_multipart_content_type_parameters(value: &str) -> Option<Vec<&str>> {
|
||||
let mut segments = Vec::new();
|
||||
let mut start = 0usize;
|
||||
let mut in_quotes = false;
|
||||
let mut escaped = false;
|
||||
|
||||
for (index, character) in value.char_indices() {
|
||||
if character.is_ascii_control() {
|
||||
return None;
|
||||
}
|
||||
if in_quotes {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
} else if character == '\\' {
|
||||
escaped = true;
|
||||
} else if character == '"' {
|
||||
in_quotes = false;
|
||||
}
|
||||
} else if character == '"' {
|
||||
in_quotes = true;
|
||||
} else if character == ';' {
|
||||
segments.push(&value[start..index]);
|
||||
start = index + character.len_utf8();
|
||||
}
|
||||
}
|
||||
|
||||
if in_quotes || escaped {
|
||||
return None;
|
||||
}
|
||||
segments.push(&value[start..]);
|
||||
Some(segments)
|
||||
}
|
||||
|
||||
fn parse_multipart_content_type_parameter_value(value: &str) -> Option<(String, bool)> {
|
||||
if value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if value.starts_with('"') {
|
||||
if value.len() < 2 || !value.ends_with('"') {
|
||||
return None;
|
||||
}
|
||||
let inner = &value[1..value.len() - 1];
|
||||
let mut parsed = String::with_capacity(inner.len());
|
||||
let mut escaped = false;
|
||||
let mut had_escape = false;
|
||||
for character in inner.chars() {
|
||||
if escaped {
|
||||
if character.is_ascii_control() {
|
||||
return None;
|
||||
}
|
||||
parsed.push(character);
|
||||
escaped = false;
|
||||
had_escape = true;
|
||||
} else if character == '\\' {
|
||||
escaped = true;
|
||||
} else {
|
||||
if character == '"' || character.is_ascii_control() {
|
||||
return None;
|
||||
}
|
||||
parsed.push(character);
|
||||
}
|
||||
}
|
||||
if escaped {
|
||||
return None;
|
||||
}
|
||||
return Some((parsed, had_escape));
|
||||
}
|
||||
|
||||
value
|
||||
.as_bytes()
|
||||
.iter()
|
||||
.copied()
|
||||
.all(is_http_token_byte)
|
||||
.then(|| (value.to_string(), false))
|
||||
}
|
||||
|
||||
fn find_multipart_boundary(haystack: &[u8], delimiter: &[u8], allow_start: bool) -> Option<usize> {
|
||||
if delimiter.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut search_start = 0usize;
|
||||
while search_start <= haystack.len() {
|
||||
let relative = find_subslice(&haystack[search_start..], delimiter)?;
|
||||
let index = search_start + relative;
|
||||
let at_line_start = index == 0
|
||||
|| (index >= 2
|
||||
&& haystack
|
||||
.get(index - 2..index)
|
||||
.is_some_and(|prefix| prefix == b"\r\n"));
|
||||
let allowed_position = if allow_start {
|
||||
at_line_start
|
||||
} else {
|
||||
index >= 2
|
||||
&& haystack
|
||||
.get(index - 2..index)
|
||||
.is_some_and(|prefix| prefix == b"\r\n")
|
||||
};
|
||||
if allowed_position && multipart_boundary_suffix_is_valid(haystack, index, delimiter) {
|
||||
return Some(index);
|
||||
}
|
||||
search_start = index.saturating_add(1);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn multipart_boundary_suffix_is_valid(haystack: &[u8], index: usize, delimiter: &[u8]) -> bool {
|
||||
let suffix_start = index.saturating_add(delimiter.len());
|
||||
let Some(suffix) = haystack.get(suffix_start..) else {
|
||||
return false;
|
||||
};
|
||||
suffix.starts_with(b"\r\n")
|
||||
|| suffix
|
||||
.strip_prefix(b"--")
|
||||
.is_some_and(|remaining| remaining.is_empty() || remaining.starts_with(b"\r\n"))
|
||||
}
|
||||
|
||||
const MAX_MULTIPART_BOUNDARY_BYTES: usize = 70;
|
||||
|
||||
fn is_valid_multipart_boundary(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= MAX_MULTIPART_BOUNDARY_BYTES
|
||||
&& value.as_bytes().iter().copied().all(is_http_token_byte)
|
||||
}
|
||||
|
||||
fn is_http_token_byte(byte: u8) -> bool {
|
||||
matches!(
|
||||
byte,
|
||||
b'0'..=b'9'
|
||||
| b'A'..=b'Z'
|
||||
| b'a'..=b'z'
|
||||
| b'!'
|
||||
| b'#'
|
||||
| b'$'
|
||||
| b'%'
|
||||
| b'&'
|
||||
| b'\''
|
||||
| b'*'
|
||||
| b'+'
|
||||
| b'-'
|
||||
| b'.'
|
||||
| b'^'
|
||||
| b'_'
|
||||
| b'`'
|
||||
| b'|'
|
||||
| b'~'
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_multipart_field(raw: &[u8]) -> Option<MultipartField> {
|
||||
let header_end = find_subslice(raw, b"\r\n\r\n")?;
|
||||
let headers = &raw[..header_end];
|
||||
let data = raw.get(header_end + 4..)?.to_vec();
|
||||
let header_text = String::from_utf8_lossy(headers);
|
||||
let header_text = std::str::from_utf8(headers).ok()?;
|
||||
|
||||
let mut name = None;
|
||||
let mut content_type = None;
|
||||
for line in header_text.lines() {
|
||||
let trimmed = line.trim();
|
||||
let lower = trimmed.to_ascii_lowercase();
|
||||
if lower.starts_with("content-disposition:") {
|
||||
name = extract_quoted_header_value(trimmed, "name");
|
||||
} else if lower.starts_with("content-type:") {
|
||||
content_type = trimmed
|
||||
.split_once(':')
|
||||
.map(|(_, value)| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
let mut disposition_seen = false;
|
||||
let mut content_type_seen = false;
|
||||
for line in header_text.split("\r\n") {
|
||||
let (header_name, header_value) = line.split_once(':')?;
|
||||
let header_name = header_name.trim();
|
||||
let header_value = header_value.trim();
|
||||
if header_name.eq_ignore_ascii_case("content-disposition") {
|
||||
if disposition_seen {
|
||||
return None;
|
||||
}
|
||||
disposition_seen = true;
|
||||
name = parse_multipart_content_disposition_name(header_value);
|
||||
} else if header_name.eq_ignore_ascii_case("content-type") {
|
||||
if content_type_seen || header_value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
content_type_seen = true;
|
||||
content_type = Some(header_value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1726,12 +1949,114 @@ fn parse_multipart_field(raw: &[u8]) -> Option<MultipartField> {
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_quoted_header_value(header: &str, key: &str) -> Option<String> {
|
||||
let pattern = format!("{key}=\"");
|
||||
let start = header.find(&pattern)? + pattern.len();
|
||||
let rest = &header[start..];
|
||||
let end = rest.find('"')?;
|
||||
Some(rest[..end].to_string())
|
||||
fn parse_multipart_content_disposition_name(value: &str) -> Option<String> {
|
||||
let segments = split_multipart_header_parameters(value)?;
|
||||
let disposition = segments.first()?.trim();
|
||||
if !disposition.eq_ignore_ascii_case("form-data") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut seen_keys = Vec::new();
|
||||
let mut name = None;
|
||||
for segment in segments.into_iter().skip(1) {
|
||||
let segment = segment.trim();
|
||||
if segment.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let (raw_key, raw_value) = segment.split_once('=')?;
|
||||
let key = raw_key.trim();
|
||||
if key.is_empty() || !key.as_bytes().iter().copied().all(is_http_token_byte) {
|
||||
return None;
|
||||
}
|
||||
if seen_keys
|
||||
.iter()
|
||||
.any(|seen: &String| seen.eq_ignore_ascii_case(key))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
seen_keys.push(key.to_ascii_lowercase());
|
||||
|
||||
let parsed_value = parse_multipart_header_parameter_value(raw_value.trim())?;
|
||||
if key.eq_ignore_ascii_case("name") {
|
||||
if parsed_value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
name = Some(parsed_value);
|
||||
}
|
||||
}
|
||||
|
||||
name
|
||||
}
|
||||
|
||||
fn split_multipart_header_parameters(value: &str) -> Option<Vec<&str>> {
|
||||
let mut segments = Vec::new();
|
||||
let mut start = 0usize;
|
||||
let mut in_quotes = false;
|
||||
let mut escaped = false;
|
||||
|
||||
for (index, byte) in value.as_bytes().iter().copied().enumerate() {
|
||||
if in_quotes {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
} else if byte == b'\\' {
|
||||
escaped = true;
|
||||
} else if byte == b'"' {
|
||||
in_quotes = false;
|
||||
}
|
||||
} else if byte == b'"' {
|
||||
in_quotes = true;
|
||||
} else if byte == b';' {
|
||||
segments.push(&value[start..index]);
|
||||
start = index + 1;
|
||||
}
|
||||
}
|
||||
|
||||
if in_quotes || escaped {
|
||||
return None;
|
||||
}
|
||||
segments.push(&value[start..]);
|
||||
Some(segments)
|
||||
}
|
||||
|
||||
fn parse_multipart_header_parameter_value(value: &str) -> Option<String> {
|
||||
if value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if value.starts_with('"') {
|
||||
if value.len() < 2 || !value.ends_with('"') {
|
||||
return None;
|
||||
}
|
||||
let inner = &value[1..value.len() - 1];
|
||||
let mut parsed = String::with_capacity(inner.len());
|
||||
let mut escaped = false;
|
||||
for character in inner.chars() {
|
||||
if escaped {
|
||||
if character.is_control() {
|
||||
return None;
|
||||
}
|
||||
parsed.push(character);
|
||||
escaped = false;
|
||||
} else if character == '\\' {
|
||||
escaped = true;
|
||||
} else {
|
||||
if character == '"' || character.is_control() {
|
||||
return None;
|
||||
}
|
||||
parsed.push(character);
|
||||
}
|
||||
}
|
||||
if escaped {
|
||||
return None;
|
||||
}
|
||||
return Some(parsed);
|
||||
}
|
||||
|
||||
value
|
||||
.as_bytes()
|
||||
.iter()
|
||||
.copied()
|
||||
.all(is_http_token_byte)
|
||||
.then(|| value.to_string())
|
||||
}
|
||||
|
||||
fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
||||
@@ -1752,10 +2077,13 @@ mod tests {
|
||||
use super::{
|
||||
build_chatgpt_web_image_request_body, build_codex_openai_image_api_provider_request_body,
|
||||
build_openai_image_api_provider_request_body, build_openai_image_provider_request_body,
|
||||
is_openai_image_stream_request, normalize_openai_image_quality,
|
||||
decode_multipart_body_base64_with_limit, find_multipart_boundary,
|
||||
is_openai_image_stream_request, multipart_boundary, normalize_openai_image_quality,
|
||||
normalize_openai_image_request, normalize_openai_image_request_with_options,
|
||||
openai_image_operation_from_path, project_codex_openai_image_api_request_body,
|
||||
project_openai_image_api_request_body, OpenAiImageNormalizeOptions, OpenAiImageOperation,
|
||||
openai_image_operation_from_path, parse_multipart_fields,
|
||||
project_codex_openai_image_api_request_body, project_openai_image_api_request_body,
|
||||
OpenAiImageNormalizeOptions, OpenAiImageOperation, MAX_MULTIPART_BOUNDARY_BYTES,
|
||||
MAX_MULTIPART_PARTS, MAX_MULTIPART_PART_HEADER_BYTES,
|
||||
};
|
||||
use crate::formats::openai::image::spec::{resolve_stream_spec, resolve_sync_spec};
|
||||
|
||||
@@ -1833,6 +2161,210 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_boundary_requires_rfc_token_and_length_limits() {
|
||||
assert_eq!(
|
||||
multipart_boundary("Multipart/Form-Data; boundary=quoted-token-123").as_deref(),
|
||||
Some("quoted-token-123")
|
||||
);
|
||||
assert_eq!(
|
||||
multipart_boundary("multipart/form-data; boundary=\"quoted-token-123\"").as_deref(),
|
||||
Some("quoted-token-123")
|
||||
);
|
||||
|
||||
for content_type in [
|
||||
"multipart/form-data; boundary=bad boundary",
|
||||
"multipart/form-data; boundary=bad\"quote",
|
||||
"multipart/form-data; boundary=\"unterminated",
|
||||
"multipart/form-data; boundary=first; boundary=second",
|
||||
"multipart/form-data; foo",
|
||||
"multipart/form-data; foo=\"unterminated; boundary=valid",
|
||||
"multipart/form-data; boundary=valid trailing",
|
||||
"application/json; boundary=valid-token",
|
||||
] {
|
||||
assert!(
|
||||
multipart_boundary(content_type).is_none(),
|
||||
"{content_type:?}"
|
||||
);
|
||||
}
|
||||
|
||||
let oversized = "a".repeat(MAX_MULTIPART_BOUNDARY_BYTES + 1);
|
||||
assert!(
|
||||
multipart_boundary(&format!("multipart/form-data; boundary={oversized}")).is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_boundary_accepts_quoted_unknown_parameter_with_semicolon() {
|
||||
assert_eq!(
|
||||
multipart_boundary("multipart/form-data; note=\"semi;colon\"; boundary=quoted-token")
|
||||
.as_deref(),
|
||||
Some("quoted-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_boundary_rejects_escaped_or_duplicate_content_type_parameters() {
|
||||
for content_type in [
|
||||
"multipart/form-data; boundary=\"escaped\\\"token\"",
|
||||
"multipart/form-data; note=one; NOTE=two; boundary=token",
|
||||
"multipart/form-data; note=\"unterminated; boundary=token",
|
||||
"multipart/form-data; note=\"closed\"trailing; boundary=token",
|
||||
] {
|
||||
assert!(
|
||||
multipart_boundary(content_type).is_none(),
|
||||
"malformed content type must be rejected: {content_type:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_parser_caps_part_count_and_header_size() {
|
||||
let boundary = "bounded-parts";
|
||||
let mut accepted_body = Vec::new();
|
||||
for index in 0..MAX_MULTIPART_PARTS {
|
||||
accepted_body.extend_from_slice(
|
||||
format!(
|
||||
"--{boundary}\r\nContent-Disposition: form-data; name=\"field-{index}\"\r\n\r\nvalue\r\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
}
|
||||
accepted_body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes());
|
||||
assert_eq!(
|
||||
parse_multipart_fields(&accepted_body, boundary).len(),
|
||||
MAX_MULTIPART_PARTS
|
||||
);
|
||||
|
||||
let mut body = Vec::new();
|
||||
for index in 0..(MAX_MULTIPART_PARTS + 1) {
|
||||
body.extend_from_slice(
|
||||
format!(
|
||||
"--{boundary}\r\nContent-Disposition: form-data; name=\"field-{index}\"\r\n\r\nvalue\r\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
}
|
||||
body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes());
|
||||
assert!(parse_multipart_fields(&body, boundary).is_empty());
|
||||
|
||||
let mut oversized_header =
|
||||
format!("--{boundary}\r\nContent-Disposition: form-data; name=\"field\"; x=\"")
|
||||
.into_bytes();
|
||||
oversized_header.extend(std::iter::repeat_n(b'x', MAX_MULTIPART_PART_HEADER_BYTES));
|
||||
oversized_header
|
||||
.extend_from_slice(format!("\"\r\n\r\nvalue\r\n--{boundary}--\r\n").as_bytes());
|
||||
assert!(parse_multipart_fields(&oversized_header, boundary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_body_base64_decode_enforces_allocation_limit() {
|
||||
let exact = vec![b'x'; 6];
|
||||
let exact_encoded = base64::engine::general_purpose::STANDARD.encode(&exact);
|
||||
assert_eq!(
|
||||
decode_multipart_body_base64_with_limit(&exact_encoded, exact.len()),
|
||||
Some(exact)
|
||||
);
|
||||
|
||||
let oversized = vec![b'x'; 7];
|
||||
let oversized_encoded = base64::engine::general_purpose::STANDARD.encode(oversized);
|
||||
assert_eq!(
|
||||
decode_multipart_body_base64_with_limit(&oversized_encoded, 6),
|
||||
None
|
||||
);
|
||||
|
||||
let same_encoded_bucket = base64::engine::general_purpose::STANDARD.encode([b'x'; 6]);
|
||||
assert_eq!(
|
||||
decode_multipart_body_base64_with_limit(&same_encoded_bucket, 5),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_parser_preserves_boundary_like_payload_and_fails_closed() {
|
||||
let boundary = "payload-boundary";
|
||||
let body = format!(
|
||||
concat!(
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; name=\"prompt\"\r\n\r\n",
|
||||
"prefix\r\n--{boundary}X\r\nsuffix--{boundary}\r\n",
|
||||
"--{boundary}--\r\n"
|
||||
),
|
||||
boundary = boundary,
|
||||
);
|
||||
let fields = parse_multipart_fields(body.as_bytes(), boundary);
|
||||
assert_eq!(fields.len(), 1);
|
||||
assert_eq!(fields[0].name, "prompt");
|
||||
assert_eq!(
|
||||
fields[0].data,
|
||||
format!("prefix\r\n--{boundary}X\r\nsuffix--{boundary}").into_bytes()
|
||||
);
|
||||
|
||||
let malformed = format!(
|
||||
concat!(
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; name=\"first\"\r\n\r\n",
|
||||
"ok\r\n",
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; name=\"second\"\r\n\r\n",
|
||||
"truncated\r\n--{boundary}X\r\n"
|
||||
),
|
||||
boundary = boundary,
|
||||
);
|
||||
assert!(parse_multipart_fields(malformed.as_bytes(), boundary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_parser_does_not_extract_name_from_filename_and_rejects_duplicates() {
|
||||
let boundary = "header-parameters";
|
||||
let filename_only = format!(
|
||||
concat!(
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; filename=\"name=\\\"prompt\\\"\"\r\n\r\n",
|
||||
"attacker-value\r\n",
|
||||
"--{boundary}--\r\n"
|
||||
),
|
||||
boundary = boundary,
|
||||
);
|
||||
assert!(parse_multipart_fields(filename_only.as_bytes(), boundary).is_empty());
|
||||
|
||||
let duplicate_name = format!(
|
||||
concat!(
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; name=\"prompt\"; name=\"image\"\r\n\r\n",
|
||||
"ambiguous-value\r\n",
|
||||
"--{boundary}--\r\n"
|
||||
),
|
||||
boundary = boundary,
|
||||
);
|
||||
assert!(parse_multipart_fields(duplicate_name.as_bytes(), boundary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_parser_rejects_garbage_after_closing_boundary() {
|
||||
let boundary = "closing-suffix";
|
||||
let body = format!(
|
||||
concat!(
|
||||
"--{boundary}\r\n",
|
||||
"Content-Disposition: form-data; name=\"prompt\"\r\n\r\n",
|
||||
"value\r\n",
|
||||
"--{boundary}--junk"
|
||||
),
|
||||
boundary = boundary,
|
||||
);
|
||||
assert!(parse_multipart_fields(body.as_bytes(), boundary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_scanner_skips_short_prefix_before_later_valid_boundary() {
|
||||
let delimiter = b"--scanner-boundary";
|
||||
let haystack = b"x--scanner-boundary\r\ncontent\r\n--scanner-boundary\r\n";
|
||||
assert_eq!(
|
||||
find_multipart_boundary(haystack, delimiter, false),
|
||||
Some(haystack.len() - delimiter.len() - 2)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_image_variation_path_is_not_supported() {
|
||||
let boundary = "boundary-variation-123";
|
||||
|
||||
@@ -1,17 +1,33 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use base64::Engine as _;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest as _, Sha256};
|
||||
|
||||
use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
|
||||
use crate::formats::openai::image::{
|
||||
bounded_openai_image_revised_prompt, is_safe_openai_image_base64_payload,
|
||||
normalize_openai_image_output_format,
|
||||
};
|
||||
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
|
||||
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
|
||||
use crate::formats::shared::stream_core::common::{
|
||||
build_openai_chat_chunk, build_openai_chat_finish_chunk,
|
||||
build_openai_chat_usage_chunk_with_cache,
|
||||
};
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
use crate::formats::shared::{decode_sync_report_body_base64, AiSurfaceFinalizeError};
|
||||
|
||||
// Bound parser carry state while still allowing the largest supported image
|
||||
// records. A 3840x2160 RGBA image is about 33 MiB before base64 encoding, so
|
||||
// 64 MiB leaves room for encoding and event metadata. This is a per-stream
|
||||
// parser limit, not a response-body or concurrency limit.
|
||||
const MAX_STREAM_REWRITE_BUFFER_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
// OpenAI image responses currently allow at most a small number of output
|
||||
// images per request. Keep enough room for valid multi-image responses while
|
||||
// preventing an untrusted provider from growing the de-duplication sets
|
||||
// without bound.
|
||||
const MAX_IMAGE_OUTPUT_KEYS: usize = 64;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct OpenAiImageStreamState {
|
||||
@@ -49,7 +65,8 @@ struct OpenAiImageChatFrame {
|
||||
#[derive(Default)]
|
||||
pub struct OpenAiImageStreamTerminalState {
|
||||
event_name: Option<String>,
|
||||
data_lines: Vec<String>,
|
||||
data: Option<String>,
|
||||
buffered_data_bytes: usize,
|
||||
response_id: Option<String>,
|
||||
model: Option<String>,
|
||||
image_count: u64,
|
||||
@@ -65,14 +82,15 @@ impl OpenAiImageStreamState {
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(block_end) = find_sse_block_end(&self.buffered) {
|
||||
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_block(report_context, &block)?);
|
||||
drain_sse_separator(&mut self.buffered);
|
||||
}
|
||||
Ok(output)
|
||||
let mut buffered = std::mem::take(&mut self.buffered);
|
||||
let result = process_bounded_sse_chunk(
|
||||
&mut buffered,
|
||||
chunk,
|
||||
MAX_STREAM_REWRITE_BUFFER_BYTES,
|
||||
|block| self.transform_block(report_context, block),
|
||||
);
|
||||
self.buffered = buffered;
|
||||
result
|
||||
}
|
||||
|
||||
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
@@ -91,16 +109,21 @@ impl OpenAiImageStreamState {
|
||||
let text = std::str::from_utf8(block)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
let mut event_name = None::<String>;
|
||||
let mut data_lines = Vec::new();
|
||||
let mut data = None::<String>;
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
let had_data = data.is_some();
|
||||
let data_value = data.get_or_insert_with(String::new);
|
||||
if had_data {
|
||||
data_value.push('\n');
|
||||
}
|
||||
data_value.push_str(value.trim());
|
||||
}
|
||||
}
|
||||
let data = data_lines.join("\n");
|
||||
let data = data.unwrap_or_default();
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -137,7 +160,7 @@ impl OpenAiImageStreamState {
|
||||
.or_else(|| event.get("b64_json"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -181,7 +204,7 @@ impl OpenAiImageStreamState {
|
||||
let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if result.is_empty() {
|
||||
if !is_safe_openai_image_base64_payload(result) {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
@@ -274,14 +297,15 @@ impl OpenAiImageChatStreamState {
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(block_end) = find_sse_block_end(&self.buffered) {
|
||||
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_block(report_context, &block)?);
|
||||
drain_sse_separator(&mut self.buffered);
|
||||
}
|
||||
Ok(output)
|
||||
let mut buffered = std::mem::take(&mut self.buffered);
|
||||
let result = process_bounded_sse_chunk(
|
||||
&mut buffered,
|
||||
chunk,
|
||||
MAX_STREAM_REWRITE_BUFFER_BYTES,
|
||||
|block| self.transform_block(report_context, block),
|
||||
);
|
||||
self.buffered = buffered;
|
||||
result
|
||||
}
|
||||
|
||||
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
@@ -305,16 +329,21 @@ impl OpenAiImageChatStreamState {
|
||||
let text = std::str::from_utf8(block)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
let mut event_name = None::<String>;
|
||||
let mut data_lines = Vec::new();
|
||||
let mut data = None::<String>;
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
let had_data = data.is_some();
|
||||
let data_value = data.get_or_insert_with(String::new);
|
||||
if had_data {
|
||||
data_value.push('\n');
|
||||
}
|
||||
data_value.push_str(value.trim());
|
||||
}
|
||||
}
|
||||
let data = data_lines.join("\n");
|
||||
let data = data.unwrap_or_default();
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -355,16 +384,15 @@ impl OpenAiImageChatStreamState {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) {
|
||||
if !result.is_empty() {
|
||||
if is_safe_openai_image_base64_payload(result) {
|
||||
let key = image_chat_output_key(item, result);
|
||||
if self.emitted_image_keys.insert(key) {
|
||||
if insert_bounded_image_key(&mut self.emitted_image_keys, key) {
|
||||
self.latest_image = Some(OpenAiImageChatFrame {
|
||||
b64_json: result.to_string(),
|
||||
output_format: item
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.map(ToOwned::to_owned),
|
||||
});
|
||||
self.emitted_image_count = self.emitted_image_count.saturating_add(1);
|
||||
@@ -417,15 +445,14 @@ impl OpenAiImageChatStreamState {
|
||||
.or_else(|| event.get("result"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
{
|
||||
self.latest_image = Some(OpenAiImageChatFrame {
|
||||
b64_json: result.to_string(),
|
||||
output_format: event
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.map(ToOwned::to_owned),
|
||||
});
|
||||
self.emitted_image_count = self.emitted_image_count.max(1);
|
||||
@@ -592,7 +619,11 @@ impl OpenAiImageStreamTerminalState {
|
||||
if let Some(value) = trimmed.strip_prefix("event:") {
|
||||
self.event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = trimmed.strip_prefix("data:") {
|
||||
self.data_lines.push(value.trim().to_string());
|
||||
append_bounded_sse_data_line(
|
||||
&mut self.data,
|
||||
&mut self.buffered_data_bytes,
|
||||
value.trim(),
|
||||
)?;
|
||||
}
|
||||
Ok(self.latest_summary(report_context))
|
||||
}
|
||||
@@ -609,11 +640,12 @@ impl OpenAiImageStreamTerminalState {
|
||||
}
|
||||
|
||||
fn flush_event(&mut self, report_context: &Value) -> Result<(), AiSurfaceFinalizeError> {
|
||||
if self.data_lines.is_empty() {
|
||||
let Some(data) = self.data.take() else {
|
||||
self.buffered_data_bytes = 0;
|
||||
self.event_name = None;
|
||||
return Ok(());
|
||||
}
|
||||
let data = std::mem::take(&mut self.data_lines).join("\n");
|
||||
};
|
||||
self.buffered_data_bytes = 0;
|
||||
let event_name = self.event_name.take();
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(());
|
||||
@@ -660,12 +692,12 @@ impl OpenAiImageStreamTerminalState {
|
||||
.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let key = image_chat_output_key(item, result);
|
||||
if self.image_keys.insert(key) {
|
||||
if insert_bounded_image_key(&mut self.image_keys, key) {
|
||||
self.image_count = self.image_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
@@ -695,7 +727,7 @@ impl OpenAiImageStreamTerminalState {
|
||||
.or_else(|| event.get("result"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
.is_some_and(is_safe_openai_image_base64_payload)
|
||||
{
|
||||
self.image_count = 1;
|
||||
}
|
||||
@@ -759,7 +791,7 @@ fn completed_response_image_chat_frame(response: &Value) -> Option<OpenAiImageCh
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.find_map(|item| {
|
||||
let result = item.get("result").and_then(Value::as_str)?.trim();
|
||||
if result.is_empty() {
|
||||
if !is_safe_openai_image_base64_payload(result) {
|
||||
return None;
|
||||
}
|
||||
Some(OpenAiImageChatFrame {
|
||||
@@ -767,8 +799,7 @@ fn completed_response_image_chat_frame(response: &Value) -> Option<OpenAiImageCh
|
||||
output_format: item
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.map(ToOwned::to_owned),
|
||||
})
|
||||
})
|
||||
@@ -785,19 +816,33 @@ fn completed_response_image_count(response: &Value) -> u64 {
|
||||
item.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
.is_some_and(is_safe_openai_image_base64_payload)
|
||||
})
|
||||
.take(MAX_IMAGE_OUTPUT_KEYS)
|
||||
.count() as u64
|
||||
}
|
||||
|
||||
fn image_chat_output_key(item: &Map<String, Value>, result: &str) -> String {
|
||||
item.get("id")
|
||||
let (source, value) = item
|
||||
.get("id")
|
||||
.or_else(|| item.get("call_id"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| result.to_string())
|
||||
.map(|value| ("id", value))
|
||||
.unwrap_or(("result", result.trim()));
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(source.as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update(value.as_bytes());
|
||||
format!("{source}:{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn insert_bounded_image_key(keys: &mut BTreeSet<String>, key: String) -> bool {
|
||||
if keys.contains(&key) || keys.len() >= MAX_IMAGE_OUTPUT_KEYS {
|
||||
return false;
|
||||
}
|
||||
keys.insert(key)
|
||||
}
|
||||
|
||||
fn openai_image_stream_standardized_usage(
|
||||
@@ -895,19 +940,13 @@ fn openai_image_usage_to_standardized_usage(value: &Value) -> Option<Standardize
|
||||
}
|
||||
|
||||
fn image_chat_markdown(frame: &OpenAiImageChatFrame) -> String {
|
||||
let mime_type = match frame
|
||||
.output_format
|
||||
.as_deref()
|
||||
.unwrap_or("png")
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"jpg" | "jpeg" => "image/jpeg".to_string(),
|
||||
"webp" => "image/webp".to_string(),
|
||||
"png" => "image/png".to_string(),
|
||||
value if !value.is_empty() => format!("image/{value}"),
|
||||
_ => "image/png".to_string(),
|
||||
let mime_type = match frame.output_format.as_deref().map(str::trim) {
|
||||
Some(value) if value.eq_ignore_ascii_case("jpg") || value.eq_ignore_ascii_case("jpeg") => {
|
||||
"image/jpeg"
|
||||
}
|
||||
Some(value) if value.eq_ignore_ascii_case("webp") => "image/webp",
|
||||
Some(value) if value.eq_ignore_ascii_case("png") => "image/png",
|
||||
_ => "image/png",
|
||||
};
|
||||
format!(
|
||||
"",
|
||||
@@ -1040,7 +1079,7 @@ fn completed_response_image_result(event: &Value) -> Option<&str> {
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.filter_map(|item| item.get("result").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.find(|value| !value.is_empty())
|
||||
.find(|value| is_safe_openai_image_base64_payload(value))
|
||||
}
|
||||
|
||||
fn requested_partial_images(report_context: &Value) -> u64 {
|
||||
@@ -1126,6 +1165,82 @@ fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
|
||||
})
|
||||
}
|
||||
|
||||
fn process_bounded_sse_chunk<F>(
|
||||
buffered: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
max_bytes: usize,
|
||||
mut transform: F,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError>
|
||||
where
|
||||
F: FnMut(&[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError>,
|
||||
{
|
||||
let mut remaining = chunk;
|
||||
let mut output = Vec::new();
|
||||
loop {
|
||||
if let Some(block_end) = find_sse_block_end(buffered) {
|
||||
let block = buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(transform(&block)?);
|
||||
drain_sse_separator(buffered);
|
||||
continue;
|
||||
}
|
||||
if remaining.is_empty() {
|
||||
break;
|
||||
}
|
||||
let available = max_bytes.saturating_sub(buffered.len());
|
||||
if available == 0 {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"image stream buffer exceeds {max_bytes} bytes"
|
||||
)));
|
||||
}
|
||||
let take = available.min(remaining.len());
|
||||
append_bounded_stream_rewrite_chunk(buffered, &remaining[..take], max_bytes)?;
|
||||
remaining = &remaining[take..];
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn append_bounded_stream_rewrite_chunk(
|
||||
buffered: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
max_bytes: usize,
|
||||
) -> Result<(), AiSurfaceFinalizeError> {
|
||||
let next_len = buffered
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.ok_or_else(|| AiSurfaceFinalizeError::new("image stream buffer length overflow"))?;
|
||||
if next_len > max_bytes {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"image stream buffer exceeds {max_bytes} bytes"
|
||||
)));
|
||||
}
|
||||
buffered.extend_from_slice(chunk);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn append_bounded_sse_data_line(
|
||||
data: &mut Option<String>,
|
||||
buffered_data_bytes: &mut usize,
|
||||
value: &str,
|
||||
) -> Result<(), AiSurfaceFinalizeError> {
|
||||
let separator_bytes = usize::from(data.is_some());
|
||||
let next_len = buffered_data_bytes
|
||||
.checked_add(separator_bytes)
|
||||
.and_then(|length| length.checked_add(value.len()))
|
||||
.ok_or_else(|| AiSurfaceFinalizeError::new("image stream data buffer length overflow"))?;
|
||||
if next_len > MAX_STREAM_REWRITE_BUFFER_BYTES {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"image stream data buffer exceeds {MAX_STREAM_REWRITE_BUFFER_BYTES} bytes"
|
||||
)));
|
||||
}
|
||||
let data_value = data.get_or_insert_with(String::new);
|
||||
if separator_bytes != 0 {
|
||||
data_value.push('\n');
|
||||
}
|
||||
data_value.push_str(value);
|
||||
*buffered_data_bytes = next_len;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
|
||||
buffer
|
||||
.windows(2)
|
||||
@@ -1198,8 +1313,16 @@ pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
}
|
||||
if let Some(provider_body_json) = body_json {
|
||||
if openai_image_response_has_standard_data(provider_body_json) {
|
||||
let Some(client_body_json) = crate::formats::shared::image_bridge::
|
||||
build_openai_image_response_from_standard_image_response(
|
||||
provider_body_json,
|
||||
Some(report_context),
|
||||
)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
return Ok(Some(OpenAiImageSyncFinalizeProduct {
|
||||
client_body_json: provider_body_json.clone(),
|
||||
client_body_json,
|
||||
provider_body_json: provider_body_json.clone(),
|
||||
}));
|
||||
}
|
||||
@@ -1223,16 +1346,16 @@ pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.unwrap_or(CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT);
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
let text = std::str::from_utf8(&body_bytes)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
|
||||
let mut created = None;
|
||||
let mut completed_response = None;
|
||||
let mut images = Vec::new();
|
||||
let mut image_keys = BTreeSet::new();
|
||||
|
||||
for raw_block in text.split("\n\n") {
|
||||
let block = raw_block.trim();
|
||||
@@ -1271,10 +1394,30 @@ pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
let Some(result) = item.get("result").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
let result = result.trim();
|
||||
if !is_safe_openai_image_base64_payload(result)
|
||||
|| !insert_bounded_image_key(
|
||||
&mut image_keys,
|
||||
image_chat_output_key(item, result),
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let output_format = item
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.unwrap_or(default_output_format);
|
||||
let revised_prompt = item
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
.map(|value| Value::String(value.to_string()))
|
||||
.unwrap_or(Value::Null);
|
||||
images.push(serde_json::json!({
|
||||
"b64_json": result,
|
||||
"output_format": item.get("output_format").cloned().unwrap_or(Value::String(default_output_format.to_string())),
|
||||
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(Value::Null),
|
||||
"output_format": output_format,
|
||||
"revised_prompt": revised_prompt,
|
||||
}));
|
||||
}
|
||||
"response.completed" => {
|
||||
@@ -1357,10 +1500,16 @@ fn openai_image_response_has_standard_data(body_json: &Value) -> bool {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine as _;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::{maybe_build_openai_image_sync_finalize_product, OpenAiImageStreamState};
|
||||
use base64::Engine as _;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
completed_response_image_count, image_chat_output_key, insert_bounded_image_key,
|
||||
maybe_build_openai_image_sync_finalize_product, process_bounded_sse_chunk,
|
||||
OpenAiImageChatStreamState, OpenAiImageStreamState, OpenAiImageStreamTerminalState,
|
||||
};
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
@@ -1514,6 +1663,308 @@ mod tests {
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_rejects_unbounded_incomplete_sse_buffer() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let mut rewriter = OpenAiImageStreamState::default();
|
||||
let oversized = vec![b'x'; super::MAX_STREAM_REWRITE_BUFFER_BYTES + 1];
|
||||
|
||||
let error = rewriter
|
||||
.push_chunk(&report_context, &oversized)
|
||||
.expect_err("incomplete image SSE block must be bounded");
|
||||
assert!(error.0.contains("image stream buffer exceeds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_chat_stream_rejects_unbounded_incomplete_sse_buffer() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"needs_conversion": true,
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let mut rewriter = OpenAiImageChatStreamState::default();
|
||||
let oversized = vec![b'x'; super::MAX_STREAM_REWRITE_BUFFER_BYTES + 1];
|
||||
|
||||
let error = rewriter
|
||||
.push_chunk(&report_context, &oversized)
|
||||
.expect_err("incomplete image SSE block must be bounded");
|
||||
assert!(error.0.contains("image stream buffer exceeds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_terminal_observer_rejects_unbounded_data_lines() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let mut observer = OpenAiImageStreamTerminalState::default();
|
||||
let first_len = super::MAX_STREAM_REWRITE_BUFFER_BYTES / 2;
|
||||
let second_len = super::MAX_STREAM_REWRITE_BUFFER_BYTES - first_len;
|
||||
let data_line = |length: usize| {
|
||||
let mut line = Vec::with_capacity(6 + length);
|
||||
line.extend_from_slice(b"data: ");
|
||||
line.extend(std::iter::repeat_n(b'x', length));
|
||||
line.push(b'\n');
|
||||
line
|
||||
};
|
||||
|
||||
observer
|
||||
.push_line(&report_context, data_line(first_len))
|
||||
.expect("first data line should fit");
|
||||
let error = observer
|
||||
.push_line(&report_context, data_line(second_len))
|
||||
.expect_err("data lines without a blank separator must be bounded");
|
||||
assert!(error.0.contains("image stream data buffer exceeds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_terminal_observer_preserves_multiline_data_semantics() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let mut observer = OpenAiImageStreamTerminalState::default();
|
||||
|
||||
// An empty first data line must still contribute the SSE newline when
|
||||
// the following line contains the JSON event.
|
||||
observer
|
||||
.push_line(&report_context, b"data:\n".to_vec())
|
||||
.expect("empty data line should be accepted");
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_multiline\",\"model\":\"gpt-image-2\"}}\n".to_vec(),
|
||||
)
|
||||
.expect("second data line should be accepted");
|
||||
observer
|
||||
.push_line(&report_context, b"\n".to_vec())
|
||||
.expect("event separator should flush");
|
||||
|
||||
let summary = observer
|
||||
.finish(&report_context)
|
||||
.expect("observer finish should succeed")
|
||||
.expect("completed event should produce a summary");
|
||||
assert_eq!(summary.response_id.as_deref(), Some("resp_multiline"));
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-image-2"));
|
||||
assert!(summary.observed_finish);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_terminal_observer_compacts_empty_data_lines() {
|
||||
let mut observer = OpenAiImageStreamTerminalState::default();
|
||||
let mut data = None;
|
||||
let mut buffered_bytes = 0;
|
||||
for _ in 0..10_000 {
|
||||
super::append_bounded_sse_data_line(&mut data, &mut buffered_bytes, "")
|
||||
.expect("empty data line should fit");
|
||||
}
|
||||
|
||||
assert_eq!(buffered_bytes, 9_999);
|
||||
let expected = "\n".repeat(9_999);
|
||||
assert_eq!(data.as_deref(), Some(expected.as_str()));
|
||||
// Keep the state type exercised as well; this guards against changing
|
||||
// the compact representation back to per-line allocations.
|
||||
observer.data = data;
|
||||
observer.buffered_data_bytes = buffered_bytes;
|
||||
assert_eq!(observer.data.as_ref().map(String::len), Some(9_999));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_output_keys_are_hashed_to_a_fixed_length() {
|
||||
let long_id = "provider-id-".to_string() + &"x".repeat(32 * 1024);
|
||||
let item_with_id = json!({"id": long_id});
|
||||
let id_key = image_chat_output_key(
|
||||
item_with_id.as_object().expect("object item"),
|
||||
"result-that-is-ignored-when-id-is-present",
|
||||
);
|
||||
assert_eq!(id_key.len(), "id:".len() + 64);
|
||||
assert!(id_key.starts_with("id:"));
|
||||
assert!(!id_key.contains("provider-id-"));
|
||||
|
||||
let long_result = "r".repeat(32 * 1024);
|
||||
let item_without_id = json!({});
|
||||
let result_key = image_chat_output_key(
|
||||
item_without_id.as_object().expect("object item"),
|
||||
&long_result,
|
||||
);
|
||||
assert_eq!(result_key.len(), "result:".len() + 64);
|
||||
assert!(result_key.starts_with("result:"));
|
||||
assert!(!result_key.contains(&long_result));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_output_key_set_is_bounded_and_keeps_dedupe_semantics() {
|
||||
let mut keys = BTreeSet::new();
|
||||
let duplicate = "id:duplicate".to_string();
|
||||
assert!(insert_bounded_image_key(&mut keys, duplicate.clone()));
|
||||
assert!(!insert_bounded_image_key(&mut keys, duplicate.clone()));
|
||||
|
||||
for index in 1..super::MAX_IMAGE_OUTPUT_KEYS {
|
||||
assert!(insert_bounded_image_key(
|
||||
&mut keys,
|
||||
format!("id:{index:064x}"),
|
||||
));
|
||||
}
|
||||
assert_eq!(keys.len(), super::MAX_IMAGE_OUTPUT_KEYS);
|
||||
assert!(!insert_bounded_image_key(
|
||||
&mut keys,
|
||||
"id:overflow".to_string()
|
||||
));
|
||||
assert_eq!(keys.len(), super::MAX_IMAGE_OUTPUT_KEYS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_chat_markdown_does_not_embed_untrusted_output_format() {
|
||||
let frame = super::OpenAiImageChatFrame {
|
||||
b64_json: "aGVsbG8=".to_string(),
|
||||
output_format: Some("png);https://attacker.invalid/?x=(x".to_string()),
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
super::image_chat_markdown(&frame),
|
||||
""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completed_response_image_count_is_bounded() {
|
||||
let output = (0..(super::MAX_IMAGE_OUTPUT_KEYS + 8))
|
||||
.map(|index| {
|
||||
json!({
|
||||
"type": "image_generation_call",
|
||||
"result": format!("aGVsbG{index:02x}"),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let response = json!({"output": output});
|
||||
|
||||
assert_eq!(
|
||||
completed_response_image_count(&response),
|
||||
super::MAX_IMAGE_OUTPUT_KEYS as u64
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_sync_finalize_dedupes_and_bounds_output_images() {
|
||||
let report_context = json!({
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"output_format": "png"
|
||||
}
|
||||
});
|
||||
let mut stream = String::new();
|
||||
let append_output_item = |stream: &mut String, id: &str, result: &str| {
|
||||
stream.push_str("data: ");
|
||||
stream.push_str(
|
||||
&serde_json::to_string(&json!({
|
||||
"type": "response.output_item.done",
|
||||
"item": {
|
||||
"id": id,
|
||||
"type": "image_generation_call",
|
||||
"result": result,
|
||||
}
|
||||
}))
|
||||
.expect("event should serialize"),
|
||||
);
|
||||
stream.push_str("\n\n");
|
||||
};
|
||||
|
||||
append_output_item(&mut stream, "duplicate", "Zmlyc3QtaW1hZ2U=");
|
||||
append_output_item(&mut stream, "duplicate", "c2Vjb25kLWltYWdl");
|
||||
for index in 0..super::MAX_IMAGE_OUTPUT_KEYS {
|
||||
let id = format!("image-{index}");
|
||||
append_output_item(&mut stream, &id, "aGVsbG8=");
|
||||
}
|
||||
stream.push_str(
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-images\"}}\n\n",
|
||||
);
|
||||
let body_base64 = base64::engine::general_purpose::STANDARD.encode(stream.as_bytes());
|
||||
|
||||
let product = maybe_build_openai_image_sync_finalize_product(
|
||||
"openai_image_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
None,
|
||||
Some(&body_base64),
|
||||
)
|
||||
.expect("finalize should succeed")
|
||||
.expect("image stream should finalize");
|
||||
|
||||
assert_eq!(
|
||||
product.client_body_json["data"]
|
||||
.as_array()
|
||||
.expect("client data array")
|
||||
.len(),
|
||||
super::MAX_IMAGE_OUTPUT_KEYS
|
||||
);
|
||||
assert_eq!(
|
||||
product.provider_body_json["output"]
|
||||
.as_array()
|
||||
.expect("provider output array")
|
||||
.len(),
|
||||
super::MAX_IMAGE_OUTPUT_KEYS
|
||||
);
|
||||
assert_eq!(
|
||||
product.client_body_json["data"][0]["b64_json"],
|
||||
"Zmlyc3QtaW1hZ2U="
|
||||
);
|
||||
assert!(product.client_body_json["data"]
|
||||
.as_array()
|
||||
.expect("client data array")
|
||||
.iter()
|
||||
.all(|image| image["b64_json"] != "c2Vjb25kLWltYWdl"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_consumes_complete_frames_before_chunk_limit() {
|
||||
let frame = b"data: {\"type\":\"noop\"}\n\n";
|
||||
let mut chunk = Vec::with_capacity(frame.len() * 2);
|
||||
chunk.extend_from_slice(frame);
|
||||
chunk.extend_from_slice(frame);
|
||||
let mut buffered = Vec::new();
|
||||
|
||||
let output = process_bounded_sse_chunk(&mut buffered, &chunk, frame.len(), |block| {
|
||||
Ok(block.to_vec())
|
||||
})
|
||||
.expect("complete frames should be consumed even when the chunk is larger than the cap");
|
||||
|
||||
assert_eq!(output, chunk);
|
||||
assert!(buffered.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_bounds_incomplete_frame_with_small_test_limit() {
|
||||
let mut buffered = Vec::new();
|
||||
let error = process_bounded_sse_chunk(&mut buffered, b"123456789", 8, |_| Ok(Vec::new()))
|
||||
.expect_err("an incomplete frame above the cap must be rejected");
|
||||
assert!(error.0.contains("image stream buffer exceeds 8 bytes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_buffer_budget_covers_gpt_image_2_max_resolution() {
|
||||
// gpt-image-2 accepts up to 3840x2160. A worst-case raw RGBA payload
|
||||
// still fits after base64 encoding with room for SSE/JSON metadata.
|
||||
let raw_rgba_bytes = 3840usize * 2160 * 4;
|
||||
let base64_bytes = raw_rgba_bytes.div_ceil(3) * 4;
|
||||
let metadata_margin = 8 * 1024 * 1024;
|
||||
assert!(
|
||||
super::MAX_STREAM_REWRITE_BUFFER_BYTES >= base64_bytes + metadata_margin,
|
||||
"image parser cap must cover the largest supported image payload"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_finalize_product_maps_stream_response_to_client_and_provider_bodies() {
|
||||
let report_context = json!({
|
||||
@@ -1594,4 +2045,87 @@ mod tests {
|
||||
assert_eq!(product.client_body_json["data"][0]["b64_json"], "aGVsbG8=");
|
||||
assert_eq!(product.provider_body_json, provider_body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_finalize_filters_untrusted_standard_openai_image_fields() {
|
||||
let oversized_prompt = "p".repeat(256 * 1024 + 1);
|
||||
let provider_body = json!({
|
||||
"created": 1779273523,
|
||||
"model": "gpt-image-2",
|
||||
"data": [
|
||||
{"url": "javascript:alert(1)"},
|
||||
{"url": "data:text/html;base64,PGh0bWw+"},
|
||||
{
|
||||
"b64_json": "aGVsbG8=",
|
||||
"output_format": "text/html;javascript:alert(1)",
|
||||
"revised_prompt": oversized_prompt.clone()
|
||||
},
|
||||
{"b64_json": "d29ybGQ=", "output_format": "JPG"}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 2}
|
||||
});
|
||||
let report_context = json!({
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"image_request": {"operation": "generate", "output_format": "png"}
|
||||
});
|
||||
|
||||
let product = maybe_build_openai_image_sync_finalize_product(
|
||||
"openai_image_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
Some(&provider_body),
|
||||
None,
|
||||
)
|
||||
.expect("standard image response should finalize")
|
||||
.expect("at least one safe image should remain");
|
||||
|
||||
let client_data = product.client_body_json["data"]
|
||||
.as_array()
|
||||
.expect("client data array");
|
||||
assert_eq!(client_data.len(), 2);
|
||||
assert_eq!(client_data[0]["b64_json"], "aGVsbG8=");
|
||||
assert_eq!(client_data[0]["revised_prompt"], Value::Null);
|
||||
assert!(client_data[0].get("output_format").is_none());
|
||||
assert_eq!(client_data[1]["b64_json"], "d29ybGQ=");
|
||||
assert_eq!(client_data[1]["output_format"], "jpeg");
|
||||
let serialized = serde_json::to_string(&product.client_body_json).expect("json");
|
||||
assert!(!serialized.contains("javascript:"));
|
||||
assert!(!serialized.contains("text/html"));
|
||||
assert!(!serialized.contains(&oversized_prompt));
|
||||
assert_eq!(product.provider_body_json, provider_body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_finalize_accepts_maximum_safe_base64_image_payload() {
|
||||
let payload = "A".repeat(crate::formats::openai::image::MAX_OPENAI_IMAGE_DATA_BYTES);
|
||||
let provider_body = json!({
|
||||
"created": 1779273523,
|
||||
"data": [{"b64_json": payload}]
|
||||
});
|
||||
let report_context = json!({
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
|
||||
let product = maybe_build_openai_image_sync_finalize_product(
|
||||
"openai_image_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
Some(&provider_body),
|
||||
None,
|
||||
)
|
||||
.expect("maximum safe image response should finalize")
|
||||
.expect("maximum safe image should be retained");
|
||||
|
||||
let returned = product.client_body_json["data"][0]["b64_json"]
|
||||
.as_str()
|
||||
.expect("base64 payload");
|
||||
assert_eq!(
|
||||
returned.len(),
|
||||
crate::formats::openai::image::MAX_OPENAI_IMAGE_DATA_BYTES
|
||||
);
|
||||
assert!(returned.as_bytes().iter().all(|byte| *byte == b'A'));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,21 @@
|
||||
use serde_json::{json, Map, Number, Value};
|
||||
|
||||
use crate::formats::openai::image::{
|
||||
bounded_openai_image_revised_prompt, is_safe_openai_image_base64_payload,
|
||||
normalize_openai_image_output_format, parse_safe_openai_image_data_url,
|
||||
safe_openai_image_mime_type, sanitize_openai_image_source_url,
|
||||
};
|
||||
use crate::formats::shared::model_directives::extract_gemini_model_from_path;
|
||||
|
||||
const MAX_IMAGE_BRIDGE_OUTPUTS: usize = 64;
|
||||
// Responses output items are provider-controlled and may contain deeply
|
||||
// nested message/content arrays. Keep the projection bounded independently
|
||||
// of the image count so a pathological text envelope cannot exhaust stack or
|
||||
// heap while an otherwise valid image is being bridged.
|
||||
const MAX_IMAGE_BRIDGE_PARTS: usize = 512;
|
||||
const MAX_IMAGE_BRIDGE_TEXT_BYTES: usize = 256 * 1024;
|
||||
const MAX_IMAGE_BRIDGE_RECURSION_DEPTH: usize = 32;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct OpenAiImageRequestForGemini {
|
||||
pub requested_model: String,
|
||||
@@ -208,7 +222,7 @@ pub fn build_openai_image_response_from_gemini_response(
|
||||
.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
{
|
||||
revised_prompt = Some(Value::String(text.to_string()));
|
||||
}
|
||||
@@ -220,6 +234,12 @@ pub fn build_openai_image_response_from_gemini_response(
|
||||
"output_format": output_format_from_mime_type(&mime_type),
|
||||
"revised_prompt": revised_prompt.clone().unwrap_or(Value::Null),
|
||||
}));
|
||||
if images.len() >= MAX_IMAGE_BRIDGE_OUTPUTS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if images.len() >= MAX_IMAGE_BRIDGE_OUTPUTS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if images.is_empty() {
|
||||
@@ -249,17 +269,64 @@ pub fn build_openai_image_response_from_gemini_response(
|
||||
Some(Value::Object(response))
|
||||
}
|
||||
|
||||
/// Projects a native OpenAI Images response before it is returned to a client.
|
||||
///
|
||||
/// Native image responses normally do not need a format conversion, but they
|
||||
/// still cross the provider trust boundary. Keep this projection separate
|
||||
/// from the raw provider body retained for the conversion/audit report so
|
||||
/// untrusted URLs, payloads, and metadata cannot be passed through unchanged.
|
||||
pub(crate) fn build_openai_image_response_from_standard_image_response(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let data = provider_body_json.get("data")?.as_array()?;
|
||||
let images = data
|
||||
.iter()
|
||||
.take(MAX_IMAGE_BRIDGE_OUTPUTS)
|
||||
.filter_map(standard_openai_image_item_to_image_data)
|
||||
.collect::<Vec<_>>();
|
||||
if images.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let created = provider_body_json
|
||||
.get("created")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or_default();
|
||||
let mut response = Map::new();
|
||||
response.insert("created".to_string(), Value::Number(Number::from(created)));
|
||||
response.insert("data".to_string(), Value::Array(images));
|
||||
if let Some(model) = provider_body_json
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| report_context.and_then(context_model))
|
||||
{
|
||||
response.insert("model".to_string(), Value::String(model.to_string()));
|
||||
}
|
||||
if let Some(usage) = provider_body_json.get("usage") {
|
||||
response.insert("usage".to_string(), usage.clone());
|
||||
}
|
||||
Some(Value::Object(response))
|
||||
}
|
||||
|
||||
pub fn build_gemini_image_response_from_openai_image_response(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let mut parts = Vec::new();
|
||||
for item in provider_body_json.get("data")?.as_array()? {
|
||||
for item in provider_body_json
|
||||
.get("data")?
|
||||
.as_array()?
|
||||
.iter()
|
||||
.take(MAX_IMAGE_BRIDGE_OUTPUTS)
|
||||
{
|
||||
if let Some(prompt) = item
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
{
|
||||
parts.push(json!({ "text": prompt }));
|
||||
}
|
||||
@@ -311,22 +378,23 @@ pub fn build_gemini_image_response_from_openai_responses_image_response(
|
||||
) -> Option<Value> {
|
||||
let output = provider_body_json.get("output").and_then(Value::as_array)?;
|
||||
let mut parts = Vec::new();
|
||||
for item in output {
|
||||
let mut budget = GeminiImagePartBudget::default();
|
||||
for item in output.iter().take(MAX_IMAGE_BRIDGE_OUTPUTS) {
|
||||
let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default();
|
||||
if item_type == "image_generation_call" {
|
||||
if let Some(prompt) = item
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
{
|
||||
parts.push(json!({ "text": prompt }));
|
||||
budget.push_text(&mut parts, prompt);
|
||||
}
|
||||
let Some(b64_json) = item
|
||||
.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -335,19 +403,22 @@ pub fn build_gemini_image_response_from_openai_responses_image_response(
|
||||
.and_then(Value::as_str)
|
||||
.map(mime_type_from_output_format)
|
||||
.unwrap_or_else(|| "image/png".to_string());
|
||||
parts.push(json!({
|
||||
"inlineData": {
|
||||
"mimeType": mime_type,
|
||||
"data": b64_json,
|
||||
}
|
||||
}));
|
||||
budget.push_part(
|
||||
&mut parts,
|
||||
json!({
|
||||
"inlineData": {
|
||||
"mimeType": mime_type,
|
||||
"data": b64_json,
|
||||
}
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if matches!(
|
||||
item_type,
|
||||
"message" | "output_text" | "text" | "output_image" | "image_url"
|
||||
) {
|
||||
collect_openai_response_output_item_for_gemini(item, &mut parts);
|
||||
collect_openai_response_output_item_for_gemini(item, &mut parts, &mut budget, 0);
|
||||
}
|
||||
}
|
||||
if !parts.iter().any(is_gemini_inline_image_part) {
|
||||
@@ -389,6 +460,7 @@ pub fn build_openai_image_response_from_response_stream_sync_body(
|
||||
let output = provider_body_json.get("output").and_then(Value::as_array)?;
|
||||
let images = output
|
||||
.iter()
|
||||
.take(MAX_IMAGE_BRIDGE_OUTPUTS)
|
||||
.filter_map(openai_response_image_generation_item_to_image_data)
|
||||
.collect::<Vec<_>>();
|
||||
if images.is_empty() {
|
||||
@@ -438,28 +510,90 @@ fn openai_response_image_generation_item_to_image_data(item: &Value) -> Option<V
|
||||
.filter(|value| !value.is_empty());
|
||||
let mut image = Map::new();
|
||||
match result {
|
||||
Some(value) if value.starts_with("data:") => {
|
||||
Some(value) if value.trim_start().starts_with("data:") => {
|
||||
let (_, b64_json) = parse_data_url(value)?;
|
||||
image.insert("b64_json".to_string(), Value::String(b64_json));
|
||||
}
|
||||
Some(value) if value.starts_with("http://") || value.starts_with("https://") => {
|
||||
image.insert("url".to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
Some(value) => {
|
||||
image.insert("b64_json".to_string(), Value::String(value.to_string()));
|
||||
if let Some(url) = sanitize_openai_image_source_url(value) {
|
||||
if url.starts_with("data:") {
|
||||
let (_, b64_json) = parse_data_url(&url)?;
|
||||
image.insert("b64_json".to_string(), Value::String(b64_json));
|
||||
} else {
|
||||
image.insert("url".to_string(), Value::String(url));
|
||||
}
|
||||
} else if is_safe_openai_image_base64_payload(value) {
|
||||
image.insert("b64_json".to_string(), Value::String(value.to_string()));
|
||||
} else {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
let url = url?;
|
||||
if let Some((_, b64_json)) = parse_data_url(url) {
|
||||
let url = sanitize_openai_image_source_url(url)?;
|
||||
if let Some((_, b64_json)) = parse_data_url(&url) {
|
||||
image.insert("b64_json".to_string(), Value::String(b64_json));
|
||||
} else {
|
||||
image.insert("url".to_string(), Value::String(url.to_string()));
|
||||
image.insert("url".to_string(), Value::String(url));
|
||||
}
|
||||
}
|
||||
}
|
||||
image.insert(
|
||||
"revised_prompt".to_string(),
|
||||
item.get("revised_prompt").cloned().unwrap_or(Value::Null),
|
||||
item.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
.map(|value| Value::String(value.to_string()))
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
Some(Value::Object(image))
|
||||
}
|
||||
|
||||
fn standard_openai_image_item_to_image_data(item: &Value) -> Option<Value> {
|
||||
let object = item.as_object()?;
|
||||
let b64_json = object
|
||||
.get("b64_json")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value));
|
||||
let url = object
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(sanitize_openai_image_source_url);
|
||||
|
||||
let mut image = Map::new();
|
||||
if let Some(b64_json) = b64_json {
|
||||
image.insert("b64_json".to_string(), Value::String(b64_json.to_string()));
|
||||
} else if let Some(url) = url {
|
||||
if let Some((_, b64_json)) = parse_data_url(&url) {
|
||||
image.insert("b64_json".to_string(), Value::String(b64_json));
|
||||
} else {
|
||||
image.insert("url".to_string(), Value::String(url));
|
||||
}
|
||||
} else {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(output_format) = object
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
{
|
||||
image.insert(
|
||||
"output_format".to_string(),
|
||||
Value::String(output_format.to_string()),
|
||||
);
|
||||
}
|
||||
image.insert(
|
||||
"revised_prompt".to_string(),
|
||||
object
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
.map(|value| Value::String(value.to_string()))
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
Some(Value::Object(image))
|
||||
}
|
||||
@@ -474,12 +608,17 @@ pub fn build_openai_image_provider_body_from_response_stream_sync_body(
|
||||
}
|
||||
let output = data
|
||||
.iter()
|
||||
.take(MAX_IMAGE_BRIDGE_OUTPUTS)
|
||||
.filter_map(|item| {
|
||||
extract_openai_image_response_item(item).map(|(mime_type, _)| {
|
||||
json!({
|
||||
"type": "image_generation_call",
|
||||
"output_format": output_format_from_mime_type(&mime_type),
|
||||
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(Value::Null),
|
||||
"revised_prompt": item.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
.map(|value| Value::String(value.to_string()))
|
||||
.unwrap_or(Value::Null),
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -578,7 +717,8 @@ fn openai_input_image_to_gemini_part(image: Value) -> Option<Value> {
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
if let Some((mime_type, data)) = parse_data_url(image_url) {
|
||||
let image_url = sanitize_openai_image_source_url(image_url)?;
|
||||
if let Some((mime_type, data)) = parse_data_url(&image_url) {
|
||||
return Some(json!({
|
||||
"inlineData": {
|
||||
"mimeType": mime_type,
|
||||
@@ -588,7 +728,7 @@ fn openai_input_image_to_gemini_part(image: Value) -> Option<Value> {
|
||||
}
|
||||
Some(json!({
|
||||
"fileData": {
|
||||
"mimeType": mime_type_from_url(image_url),
|
||||
"mimeType": mime_type_from_url(&image_url),
|
||||
"fileUri": image_url,
|
||||
}
|
||||
}))
|
||||
@@ -691,13 +831,53 @@ fn collect_gemini_part(part: &Value, text: &mut Vec<String>, content: &mut Vec<V
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_openai_response_output_item_for_gemini(item: &Value, parts: &mut Vec<Value>) {
|
||||
#[derive(Default)]
|
||||
struct GeminiImagePartBudget {
|
||||
text_bytes: usize,
|
||||
}
|
||||
|
||||
impl GeminiImagePartBudget {
|
||||
fn push_part(&mut self, parts: &mut Vec<Value>, part: Value) {
|
||||
if parts.len() < MAX_IMAGE_BRIDGE_PARTS {
|
||||
parts.push(part);
|
||||
}
|
||||
}
|
||||
|
||||
fn push_text(&mut self, parts: &mut Vec<Value>, text: &str) {
|
||||
let text_bytes = text.len();
|
||||
let Some(next_text_bytes) = self.text_bytes.checked_add(text_bytes) else {
|
||||
return;
|
||||
};
|
||||
if next_text_bytes > MAX_IMAGE_BRIDGE_TEXT_BYTES || parts.len() >= MAX_IMAGE_BRIDGE_PARTS {
|
||||
return;
|
||||
}
|
||||
parts.push(json!({ "text": text }));
|
||||
self.text_bytes = next_text_bytes;
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_openai_response_output_item_for_gemini(
|
||||
item: &Value,
|
||||
parts: &mut Vec<Value>,
|
||||
budget: &mut GeminiImagePartBudget,
|
||||
depth: usize,
|
||||
) {
|
||||
if depth >= MAX_IMAGE_BRIDGE_RECURSION_DEPTH {
|
||||
return;
|
||||
}
|
||||
if let Value::Array(items) = item {
|
||||
for child in items {
|
||||
collect_openai_response_output_item_for_gemini(child, parts, budget, depth + 1);
|
||||
if parts.len() >= MAX_IMAGE_BRIDGE_PARTS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default();
|
||||
if item_type == "message" {
|
||||
if let Some(content) = item.get("content").and_then(Value::as_array) {
|
||||
for part in content {
|
||||
collect_openai_response_output_item_for_gemini(part, parts);
|
||||
}
|
||||
if let Some(content) = item.get("content") {
|
||||
collect_openai_response_output_item_for_gemini(content, parts, budget, depth + 1);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -708,7 +888,7 @@ fn collect_openai_response_output_item_for_gemini(item: &Value, parts: &mut Vec<
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
parts.push(json!({ "text": text }));
|
||||
budget.push_text(parts, text);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -727,12 +907,15 @@ fn collect_openai_response_output_item_for_gemini(item: &Value, parts: &mut Vec<
|
||||
.filter(|value| !value.is_empty());
|
||||
if let Some(image_url) = image_url {
|
||||
if let Some((mime_type, data)) = parse_data_url(image_url) {
|
||||
parts.push(json!({
|
||||
"inlineData": {
|
||||
"mimeType": mime_type,
|
||||
"data": data,
|
||||
}
|
||||
}));
|
||||
budget.push_part(
|
||||
parts,
|
||||
json!({
|
||||
"inlineData": {
|
||||
"mimeType": mime_type,
|
||||
"data": data,
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -745,15 +928,14 @@ fn extract_gemini_inline_image(part: &Value) -> Option<(String, String)> {
|
||||
.get("mimeType")
|
||||
.or_else(|| object.get("mime_type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| value.starts_with("image/"))
|
||||
.and_then(safe_openai_image_mime_type)
|
||||
.unwrap_or("image/png")
|
||||
.to_string();
|
||||
let data = object
|
||||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))?
|
||||
.to_string();
|
||||
Some((mime_type, data))
|
||||
}
|
||||
@@ -764,13 +946,12 @@ fn extract_openai_image_response_item(item: &Value) -> Option<(String, String)>
|
||||
.get("b64_json")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
{
|
||||
let output_format = object
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.unwrap_or("png");
|
||||
return Some((
|
||||
mime_type_from_output_format(output_format),
|
||||
@@ -786,13 +967,7 @@ fn extract_openai_image_response_item(item: &Value) -> Option<(String, String)>
|
||||
}
|
||||
|
||||
fn parse_data_url(value: &str) -> Option<(String, String)> {
|
||||
let (metadata, payload) = value.trim().split_once(',')?;
|
||||
let metadata = metadata.strip_prefix("data:")?;
|
||||
let mime_type = metadata.strip_suffix(";base64")?;
|
||||
let payload = payload.trim();
|
||||
if payload.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let (mime_type, payload) = parse_safe_openai_image_data_url(value)?;
|
||||
Some((mime_type.to_string(), payload.to_string()))
|
||||
}
|
||||
|
||||
@@ -822,12 +997,11 @@ fn output_format_from_mime_type(mime_type: &str) -> &'static str {
|
||||
}
|
||||
|
||||
fn mime_type_from_output_format(output_format: &str) -> String {
|
||||
match output_format.trim().to_ascii_lowercase().as_str() {
|
||||
"jpeg" | "jpg" => "image/jpeg".to_string(),
|
||||
"webp" => "image/webp".to_string(),
|
||||
"png" => "image/png".to_string(),
|
||||
other if other.starts_with("image/") => other.to_string(),
|
||||
_ => "image/png".to_string(),
|
||||
match normalize_openai_image_output_format(output_format) {
|
||||
Some("jpeg") => "image/jpeg".to_string(),
|
||||
Some("webp") => "image/webp".to_string(),
|
||||
Some("png") | None => "image/png".to_string(),
|
||||
Some(_) => "image/png".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -901,7 +1075,7 @@ fn context_model(context: &Value) -> Option<&str> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use http::{Method, Request};
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_gemini_image_request_body_from_openai_image_request,
|
||||
@@ -909,6 +1083,7 @@ mod tests {
|
||||
build_openai_image_request_body_from_gemini_image_request,
|
||||
build_openai_image_response_from_gemini_response,
|
||||
build_openai_image_response_from_response_stream_sync_body,
|
||||
build_openai_image_response_from_standard_image_response,
|
||||
gemini_request_is_image_generation,
|
||||
};
|
||||
use crate::formats::openai::image::request::normalize_openai_image_request;
|
||||
@@ -1129,4 +1304,88 @@ mod tests {
|
||||
);
|
||||
assert_eq!(converted["usageMetadata"]["totalTokenCount"], 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_openai_image_bridge_filters_fields_and_bounds_outputs() {
|
||||
let provider_body = json!({
|
||||
"created": 1779273523,
|
||||
"model": "gpt-image-2",
|
||||
"data": [
|
||||
{"url": "javascript:alert(1)"},
|
||||
{"url": "data:text/html;base64,PGh0bWw+"},
|
||||
{
|
||||
"b64_json": "aGVsbG8=",
|
||||
"output_format": "text/html",
|
||||
"revised_prompt": "p".repeat(256 * 1024 + 1)
|
||||
}
|
||||
]
|
||||
});
|
||||
let converted =
|
||||
build_openai_image_response_from_standard_image_response(&provider_body, None)
|
||||
.expect("valid standard image should remain");
|
||||
assert_eq!(converted["data"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(converted["data"][0]["b64_json"], "aGVsbG8=");
|
||||
assert_eq!(converted["data"][0]["revised_prompt"], Value::Null);
|
||||
assert!(converted["data"][0].get("output_format").is_none());
|
||||
let serialized = serde_json::to_string(&converted).expect("json");
|
||||
assert!(!serialized.contains("javascript:"));
|
||||
assert!(!serialized.contains("text/html"));
|
||||
|
||||
let outputs = (0..80)
|
||||
.map(|index| json!({"b64_json": format!("image{index:02}=")}))
|
||||
.collect::<Vec<_>>();
|
||||
let bounded = build_openai_image_response_from_standard_image_response(
|
||||
&json!({"data": outputs}),
|
||||
None,
|
||||
)
|
||||
.expect("bounded standard image response should convert");
|
||||
assert_eq!(bounded["data"].as_array().map(Vec::len), Some(64));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_image_bridge_bounds_nested_text_without_losing_image() {
|
||||
let mut nested = json!({"type": "output_text", "text": "nested"});
|
||||
for _ in 0..128 {
|
||||
nested = json!({"type": "message", "content": [nested]});
|
||||
}
|
||||
let converted = super::build_gemini_image_response_from_openai_responses_image_response(
|
||||
&json!({
|
||||
"output": [
|
||||
nested,
|
||||
{"type": "image_generation_call", "result": "aGVsbG8="}
|
||||
]
|
||||
}),
|
||||
None,
|
||||
)
|
||||
.expect("nested response should still retain the image");
|
||||
let parts = converted["candidates"][0]["content"]["parts"]
|
||||
.as_array()
|
||||
.expect("gemini parts");
|
||||
assert!(parts.iter().any(super::is_gemini_inline_image_part));
|
||||
assert!(parts.len() <= super::MAX_IMAGE_BRIDGE_PARTS);
|
||||
|
||||
let large_text = "t".repeat(super::MAX_IMAGE_BRIDGE_TEXT_BYTES / 2);
|
||||
let output = (0..4)
|
||||
.map(|_| json!({"type": "output_text", "text": large_text.clone()}))
|
||||
.chain(std::iter::once(json!({
|
||||
"type": "image_generation_call",
|
||||
"result": "aGVsbG8="
|
||||
})))
|
||||
.collect::<Vec<_>>();
|
||||
let bounded = super::build_gemini_image_response_from_openai_responses_image_response(
|
||||
&json!({"output": output}),
|
||||
None,
|
||||
)
|
||||
.expect("text budget should not suppress a valid image");
|
||||
let parts = bounded["candidates"][0]["content"]["parts"]
|
||||
.as_array()
|
||||
.expect("bounded parts");
|
||||
assert!(parts.iter().any(super::is_gemini_inline_image_part));
|
||||
let text_bytes = parts
|
||||
.iter()
|
||||
.filter_map(|part| part.get("text").and_then(Value::as_str))
|
||||
.map(str::len)
|
||||
.sum::<usize>();
|
||||
assert!(text_bytes <= super::MAX_IMAGE_BRIDGE_TEXT_BYTES);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
use base64::Engine as _;
|
||||
use std::fmt;
|
||||
|
||||
/// Report bodies are produced from execution responses, whose normal decoded
|
||||
/// transport limit is 64 MiB. Keep format finalizers on the same boundary so
|
||||
/// a base64 field cannot trigger an unchecked allocation before parsing.
|
||||
pub(crate) const MAX_SYNC_REPORT_BODY_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
pub mod error_body;
|
||||
pub mod family;
|
||||
pub mod image_bridge;
|
||||
@@ -18,6 +24,35 @@ pub mod sync_products;
|
||||
pub mod sync_to_stream;
|
||||
pub mod video;
|
||||
|
||||
pub(crate) fn decode_sync_report_body_base64(
|
||||
body_base64: &str,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if body_base64.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let max_encoded_len = MAX_SYNC_REPORT_BODY_BYTES
|
||||
.checked_add(2)
|
||||
.and_then(|value| value.checked_div(3))
|
||||
.and_then(|value| value.checked_mul(4))
|
||||
.unwrap_or(usize::MAX);
|
||||
if body_base64.len() > max_encoded_len {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"sync report body exceeds {} decoded bytes",
|
||||
MAX_SYNC_REPORT_BODY_BYTES
|
||||
)));
|
||||
}
|
||||
|
||||
let bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
if bytes.len() > MAX_SYNC_REPORT_BODY_BYTES {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"sync report body exceeds {} decoded bytes",
|
||||
MAX_SYNC_REPORT_BODY_BYTES
|
||||
)));
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub use self::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
|
||||
pub use self::stream_core::{CanonicalStreamEvent, CanonicalStreamFrame};
|
||||
pub use self::stream_rewrite::{
|
||||
|
||||
@@ -441,7 +441,19 @@ pub fn sanitize_request_path(path: &str) -> Option<String> {
|
||||
// including for malformed routes that will later be rejected.
|
||||
return Some("/v1/live/{call_id}".to_string());
|
||||
}
|
||||
Some(path.to_string())
|
||||
Some(sanitize_sensitive_request_path(path))
|
||||
}
|
||||
|
||||
fn sanitize_sensitive_request_path(path: &str) -> String {
|
||||
for prefix in ["/install-tunnel/", "/install/", "/i/"] {
|
||||
if path
|
||||
.strip_prefix(prefix)
|
||||
.is_some_and(|secret| !secret.is_empty())
|
||||
{
|
||||
return format!("{prefix}[redacted]");
|
||||
}
|
||||
}
|
||||
path.to_string()
|
||||
}
|
||||
|
||||
pub fn sanitize_request_query_string(query: &str) -> Option<String> {
|
||||
@@ -1022,6 +1034,29 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_path_metadata_sanitizer_redacts_install_session_codes() {
|
||||
for (raw, expected) in [
|
||||
("/install/secret-code", "/install/[redacted]"),
|
||||
("/install/secret-code.ps1", "/install/[redacted]"),
|
||||
("/i/secret-code", "/i/[redacted]"),
|
||||
(
|
||||
"/install-tunnel/secret-code.ps1?token=also-secret",
|
||||
"/install-tunnel/[redacted]",
|
||||
),
|
||||
] {
|
||||
assert_eq!(sanitize_request_path(raw).as_deref(), Some(expected));
|
||||
assert_eq!(
|
||||
sanitize_request_path_and_query(raw, None).as_deref(),
|
||||
Some(expected)
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
sanitize_request_path("/install/").as_deref(),
|
||||
Some("/install/")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_matching_requires_openai_stream_flag() {
|
||||
assert!(!is_matching_stream_request(
|
||||
|
||||
@@ -19,6 +19,12 @@ use crate::provider_compat::surfaces::{
|
||||
provider_adaptation_should_unwrap_stream_envelope, KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
|
||||
// A stream record can legitimately contain a large tool payload, but a peer
|
||||
// must not be able to keep the rewriter allocating forever by withholding the
|
||||
// record separator. This is a parser carry-buffer bound, not a response-body
|
||||
// or stream-concurrency limit.
|
||||
const MAX_STREAM_REWRITE_BUFFER_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FinalizeStreamRewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
@@ -294,7 +300,7 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
| AiSurfaceStreamRewriteState::ModelDirectiveDisplay
|
||||
| AiSurfaceStreamRewriteState::OpenAiResponsesCompat
|
||||
| AiSurfaceStreamRewriteState::Standard(_) => {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
append_bounded_stream_rewrite_chunk(&mut self.buffered, chunk)?;
|
||||
let mut output = Vec::new();
|
||||
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
|
||||
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
|
||||
@@ -399,7 +405,7 @@ impl ClaudeReadToolStreamSanitizer {
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
append_bounded_stream_rewrite_chunk(&mut self.buffered, chunk)?;
|
||||
let mut output = Vec::new();
|
||||
while let Some(record) = drain_next_sse_record(&mut self.buffered) {
|
||||
output.extend(self.transform_record(report_context, record)?);
|
||||
@@ -523,6 +529,18 @@ impl ClaudeReadToolStreamSanitizer {
|
||||
return Ok(original_record);
|
||||
}
|
||||
if let Some(partial_json) = partial_json {
|
||||
let next_len = state
|
||||
.buffered_input_json
|
||||
.len()
|
||||
.checked_add(partial_json.len())
|
||||
.ok_or_else(|| {
|
||||
AiSurfaceFinalizeError::new("stream rewrite tool input buffer length overflow")
|
||||
})?;
|
||||
if next_len > MAX_STREAM_REWRITE_BUFFER_BYTES {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"stream rewrite tool input buffer exceeds {MAX_STREAM_REWRITE_BUFFER_BYTES} bytes"
|
||||
)));
|
||||
}
|
||||
state.buffered_input_json.push_str(partial_json);
|
||||
}
|
||||
Ok(Vec::new())
|
||||
@@ -568,6 +586,23 @@ impl ClaudeReadToolStreamSanitizer {
|
||||
}
|
||||
}
|
||||
|
||||
fn append_bounded_stream_rewrite_chunk(
|
||||
buffered: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
) -> Result<(), AiSurfaceFinalizeError> {
|
||||
let next_len = buffered
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.ok_or_else(|| AiSurfaceFinalizeError::new("stream rewrite buffer length overflow"))?;
|
||||
if next_len > MAX_STREAM_REWRITE_BUFFER_BYTES {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"stream rewrite buffer exceeds {MAX_STREAM_REWRITE_BUFFER_BYTES} bytes"
|
||||
)));
|
||||
}
|
||||
buffered.extend_from_slice(chunk);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sanitize_claude_tool_input_object(block: &mut Map<String, Value>, name: &str) -> bool {
|
||||
let Some(input) = block.get("input") else {
|
||||
return false;
|
||||
@@ -1344,6 +1379,48 @@ data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
assert!(output.contains("event: content_block_stop"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_stream_bounds_read_input_json_buffer() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
"anthropic_compatibility_profile": "claude_code_legacy",
|
||||
});
|
||||
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
|
||||
.expect("same-format claude sanitizer should exist");
|
||||
rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_read_1\",\"name\":\"Read\",\"input\":{}}}\n\n",
|
||||
)
|
||||
.expect("start should rewrite");
|
||||
|
||||
let build_delta = |partial_json: &str| {
|
||||
let payload = json!({
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": partial_json,
|
||||
}
|
||||
});
|
||||
let mut record = b"event: content_block_delta\ndata: ".to_vec();
|
||||
record.extend(serde_json::to_vec(&payload).expect("delta should serialize"));
|
||||
record.extend_from_slice(b"\n\n");
|
||||
record
|
||||
};
|
||||
let first = "x".repeat(super::MAX_STREAM_REWRITE_BUFFER_BYTES / 2);
|
||||
let second = "x".repeat(super::MAX_STREAM_REWRITE_BUFFER_BYTES / 2 + 1);
|
||||
rewriter
|
||||
.push_chunk(&build_delta(&first))
|
||||
.expect("first partial JSON should fit");
|
||||
let error = rewriter
|
||||
.push_chunk(&build_delta(&second))
|
||||
.expect_err("Read input JSON must be bounded across SSE records");
|
||||
assert!(error.0.contains("tool input buffer exceeds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_stream_preserves_other_tool_empty_pages() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -23,7 +23,7 @@ use aether_ai_formats::{
|
||||
};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::AiSurfaceFinalizeError;
|
||||
use super::{decode_sync_report_body_base64, AiSurfaceFinalizeError};
|
||||
use crate::formats::claude::messages::stream::ClaudeProviderState;
|
||||
use crate::formats::gemini::generate_content::stream::GeminiProviderState;
|
||||
use crate::formats::openai::chat::stream::{OpenAIChatProviderState, OpenAIResponsesProviderState};
|
||||
@@ -79,7 +79,7 @@ pub fn maybe_build_standard_cross_format_sync_product_from_normalized_payload(
|
||||
|
||||
let (aggregated_stream_body, aggregated_stream_api_format) = match body_base64 {
|
||||
Some(body_base64) => {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
let provider_stream_event_api_format =
|
||||
provider_stream_event_api_format_for_report_context(
|
||||
report_context,
|
||||
@@ -598,7 +598,7 @@ pub fn maybe_build_embedding_cross_format_sync_product_from_normalized_payload(
|
||||
|
||||
let provider_body_json = match body_base64 {
|
||||
Some(body_base64) => {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
serde_json::from_slice::<Value>(&body_bytes).ok()
|
||||
}
|
||||
None => body_json.cloned(),
|
||||
@@ -743,7 +743,7 @@ fn maybe_build_standard_same_format_stream_sync_body(
|
||||
let Some(body_base64) = body_base64 else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
let provider_stream_event_api_format =
|
||||
provider_stream_event_api_format_for_report_context(report_context, &provider_api_format);
|
||||
let Some(mut body) = try_aggregate_standard_chat_stream_sync_response(
|
||||
@@ -892,7 +892,7 @@ fn maybe_build_openai_responses_same_family_stream_sync_body(
|
||||
let Some(body_base64) = body_base64 else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
// Same-family clients retain the authoritative terminal body verbatim, including future
|
||||
// output item fields, but unknown intermediate event types still fail closed.
|
||||
ensure_no_unknown_openai_responses_stream_events(&body_bytes, true)?;
|
||||
@@ -984,7 +984,7 @@ fn maybe_build_openai_cross_format_provider_body_from_normalized_payload(
|
||||
) -> Result<Option<OpenAiCrossFormatProviderBody>, AiSurfaceFinalizeError> {
|
||||
let aggregated_stream_body = match body_base64 {
|
||||
Some(body_base64) => {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let body_bytes = decode_sync_report_body_base64(body_base64)?;
|
||||
let normalized_provider_api_format =
|
||||
normalize_openai_responses_family_api_format(provider_api_format);
|
||||
match normalized_provider_api_format.as_str() {
|
||||
|
||||
@@ -12,6 +12,11 @@ use crate::formats::gemini::generate_content::stream::GeminiClientEmitter;
|
||||
use crate::formats::openai::chat::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
|
||||
};
|
||||
use crate::formats::openai::image::{
|
||||
bounded_openai_image_revised_prompt, is_safe_openai_image_base64_payload,
|
||||
normalize_openai_image_output_format, parse_safe_openai_image_data_url,
|
||||
sanitize_openai_image_source_url,
|
||||
};
|
||||
use crate::formats::openai::responses::history::{
|
||||
record_converted_response_history, ResponseHistoryRecord,
|
||||
};
|
||||
@@ -27,6 +32,8 @@ use crate::formats::shared::stream_core::{
|
||||
use crate::formats::shared::stream_rewrite::maybe_build_ai_surface_stream_rewriter;
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
|
||||
const MAX_OPENAI_IMAGE_OUTPUTS: usize = 64;
|
||||
|
||||
pub struct SyncToStreamBridgeOutcome {
|
||||
pub sse_body: Vec<u8>,
|
||||
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
|
||||
@@ -198,7 +205,7 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
|
||||
let Some(image) = outputs.iter().find_map(OpenAiImageOutput::b64_json) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
let image_count = outputs.len() as u64;
|
||||
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
|
||||
let event_name = openai_image_completed_event_name(report_context);
|
||||
let sse_body = encode_json_sse(
|
||||
@@ -238,7 +245,7 @@ fn maybe_bridge_openai_image_sync_json_to_chat_stream(
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
let image_count = outputs.len() as u64;
|
||||
let summary = openai_image_terminal_summary(response, report_context, image_count);
|
||||
let response_id = openai_image_bridge_response_id(response, report_context, "chatcmpl-image");
|
||||
let model = openai_image_bridge_response_model(response, report_context);
|
||||
@@ -349,7 +356,7 @@ fn maybe_bridge_openai_image_sync_json_to_responses_stream(
|
||||
}),
|
||||
)?);
|
||||
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
let image_count = outputs.len() as u64;
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(openai_image_terminal_summary(
|
||||
@@ -372,17 +379,20 @@ struct OpenAiImageOutput {
|
||||
|
||||
impl OpenAiImageOutput {
|
||||
fn b64_json(&self) -> Option<String> {
|
||||
self.b64_json
|
||||
.clone()
|
||||
.or_else(|| self.url.as_deref().and_then(extract_base64_from_data_url))
|
||||
let value = self
|
||||
.b64_json
|
||||
.as_deref()
|
||||
.or_else(|| self.url.as_deref().and_then(extract_base64_from_data_url))?;
|
||||
is_safe_openai_image_base64_payload(value).then(|| value.to_string())
|
||||
}
|
||||
|
||||
fn source_url(&self) -> Option<String> {
|
||||
self.url.clone().or_else(|| {
|
||||
let source = self.url.clone().or_else(|| {
|
||||
self.b64_json
|
||||
.as_ref()
|
||||
.map(|value| format!("data:{};base64,{value}", self.mime_type))
|
||||
})
|
||||
})?;
|
||||
sanitize_openai_image_source_url(&source)
|
||||
}
|
||||
|
||||
fn markdown(&self, index: usize) -> String {
|
||||
@@ -391,7 +401,10 @@ impl OpenAiImageOutput {
|
||||
} else {
|
||||
format!("generated image {}", index + 1)
|
||||
};
|
||||
match self.source_url() {
|
||||
match self
|
||||
.source_url()
|
||||
.and_then(|url| escape_markdown_image_destination(&url))
|
||||
{
|
||||
Some(url) => format!(""),
|
||||
None => String::new(),
|
||||
}
|
||||
@@ -408,7 +421,11 @@ impl OpenAiImageOutput {
|
||||
Value::String("image_generation_call".to_string()),
|
||||
);
|
||||
item.insert("status".to_string(), Value::String("completed".to_string()));
|
||||
if let Some(result) = self.b64_json().or_else(|| self.url.clone()) {
|
||||
if let Some(result) = self.b64_json().or_else(|| {
|
||||
self.url
|
||||
.as_deref()
|
||||
.and_then(sanitize_openai_image_source_url)
|
||||
}) {
|
||||
item.insert("result".to_string(), Value::String(result));
|
||||
}
|
||||
if let Some(output_format) = self.output_format.as_ref() {
|
||||
@@ -470,6 +487,7 @@ fn collect_openai_image_outputs(
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|item| openai_image_output_from_item(item, report_context))
|
||||
.take(MAX_OPENAI_IMAGE_OUTPUTS)
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -483,7 +501,7 @@ fn openai_image_output_from_item(
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
.and_then(sanitize_openai_image_source_url);
|
||||
if b64_json.is_none() && url.is_none() {
|
||||
return None;
|
||||
}
|
||||
@@ -491,10 +509,13 @@ fn openai_image_output_from_item(
|
||||
.get("output_format")
|
||||
.or_else(|| item.get("format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(normalize_openai_image_output_format)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_request_output_format(report_context));
|
||||
.or_else(|| {
|
||||
image_request_output_format(report_context).and_then(|value| {
|
||||
normalize_openai_image_output_format(&value).map(ToOwned::to_owned)
|
||||
})
|
||||
});
|
||||
let mime_type = url
|
||||
.as_deref()
|
||||
.and_then(extract_mime_type_from_data_url)
|
||||
@@ -507,8 +528,7 @@ fn openai_image_output_from_item(
|
||||
let revised_prompt = item
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(bounded_openai_image_revised_prompt)
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(OpenAiImageOutput {
|
||||
@@ -520,14 +540,6 @@ fn openai_image_output_from_item(
|
||||
})
|
||||
}
|
||||
|
||||
fn openai_image_response_image_count(response: &Map<String, Value>) -> u64 {
|
||||
response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| items.len() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn openai_image_terminal_summary(
|
||||
response: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
@@ -680,11 +692,12 @@ fn image_request_quality(report_context: Option<&Value>) -> Option<String> {
|
||||
}
|
||||
|
||||
fn mime_type_from_image_output_format(output_format: &str) -> String {
|
||||
match output_format.trim().to_ascii_lowercase().as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg".to_string(),
|
||||
"webp" => "image/webp".to_string(),
|
||||
"png" => "image/png".to_string(),
|
||||
value if !value.is_empty() => format!("image/{value}"),
|
||||
match output_format.trim() {
|
||||
value if value.eq_ignore_ascii_case("jpg") || value.eq_ignore_ascii_case("jpeg") => {
|
||||
"image/jpeg".to_string()
|
||||
}
|
||||
value if value.eq_ignore_ascii_case("webp") => "image/webp".to_string(),
|
||||
value if value.eq_ignore_ascii_case("png") => "image/png".to_string(),
|
||||
_ => "image/png".to_string(),
|
||||
}
|
||||
}
|
||||
@@ -913,29 +926,36 @@ fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) ->
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| is_safe_openai_image_base64_payload(value))
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| {
|
||||
item.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(extract_base64_from_data_url)
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_url(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
let (metadata, payload) = trimmed.split_once(',')?;
|
||||
if !metadata.starts_with("data:") || !metadata.ends_with(";base64") {
|
||||
return None;
|
||||
}
|
||||
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
|
||||
fn extract_base64_from_data_url(value: &str) -> Option<&str> {
|
||||
parse_safe_openai_image_data_url(value).map(|(_, payload)| payload)
|
||||
}
|
||||
|
||||
fn extract_mime_type_from_data_url(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
let (metadata, _) = trimmed.split_once(',')?;
|
||||
let mime_type = metadata.strip_prefix("data:")?.strip_suffix(";base64")?;
|
||||
let mime_type = mime_type.trim();
|
||||
(!mime_type.is_empty()).then(|| mime_type.to_string())
|
||||
parse_safe_openai_image_data_url(value).map(|(mime_type, _)| mime_type.to_string())
|
||||
}
|
||||
|
||||
fn escape_markdown_image_destination(value: &str) -> Option<String> {
|
||||
let mut escaped = String::with_capacity(value.len());
|
||||
for character in value.chars() {
|
||||
if character.is_ascii_control() || character.is_whitespace() {
|
||||
return None;
|
||||
}
|
||||
if matches!(character, '\\' | '(' | ')') {
|
||||
escaped.push('\\');
|
||||
}
|
||||
escaped.push(character);
|
||||
}
|
||||
Some(escaped)
|
||||
}
|
||||
|
||||
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
|
||||
@@ -1320,7 +1340,10 @@ fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsa
|
||||
mod tests {
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{maybe_bridge_standard_sync_json_to_stream, standardized_usage_from_openai_usage};
|
||||
use super::{
|
||||
maybe_bridge_standard_sync_json_to_stream, standardized_usage_from_openai_usage,
|
||||
OpenAiImageOutput,
|
||||
};
|
||||
use crate::formats::openai::namespace::NamespaceToolAliases;
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
@@ -1890,6 +1913,98 @@ mod tests {
|
||||
assert!(output.contains("\"total_tokens\":9"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_http_url_with_markdown_escaping() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"id": "img_url_123",
|
||||
"data": [{
|
||||
"url": "https://cdn.example.test/generated/(image).png"
|
||||
}]
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:chat",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("valid HTTP image URL should bridge");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("generated/\\\\(image\\\\).png"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_http_image_urls_and_untrusted_data_mime_types() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
for item in [
|
||||
json!({"url": "javascript:alert(1)"}),
|
||||
json!({"url": "data:text/html;base64,PGh0bWw+"}),
|
||||
] {
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({"data": [item]}),
|
||||
"openai:image",
|
||||
"openai:chat",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should not error on an unsupported image source");
|
||||
assert!(outcome.is_none(), "unsupported source must not be emitted");
|
||||
}
|
||||
|
||||
let output = OpenAiImageOutput {
|
||||
b64_json: Some("aGVsbG8=".to_string()),
|
||||
url: None,
|
||||
mime_type: "image/png".to_string(),
|
||||
output_format: Some("text/html);javascript:alert(1)".to_string()),
|
||||
revised_prompt: None,
|
||||
}
|
||||
.markdown(0);
|
||||
assert_eq!(output, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_summary_counts_only_safe_emitted_outputs() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"image_request": {"operation": "generate"}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"data": [
|
||||
{"b64_json": "not-base64!"},
|
||||
{"url": "javascript:alert(1)"},
|
||||
{"b64_json": "aGVsbG8="}
|
||||
],
|
||||
"usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:chat",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should ignore unsupported outputs")
|
||||
.expect("valid output should still bridge");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("data:image/png;base64,aGVsbG8="));
|
||||
assert!(!output.contains("not-base64"));
|
||||
assert!(!output.contains("javascript:"));
|
||||
let usage = outcome
|
||||
.terminal_summary
|
||||
.and_then(|summary| summary.standardized_usage)
|
||||
.expect("safe output should produce usage");
|
||||
assert_eq!(usage.request_count, 1);
|
||||
assert_eq!(usage.dimensions.get("image_count"), Some(&json!(1)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_aether_sse_response_capture_to_same_client_stream() {
|
||||
let captured_body = concat!(
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||
use std::fmt;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Map, Value};
|
||||
@@ -57,7 +58,7 @@ pub enum CanonicalStopReason {
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum CanonicalToolChoice {
|
||||
Auto,
|
||||
@@ -66,7 +67,7 @@ pub enum CanonicalToolChoice {
|
||||
Tool { name: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum CanonicalContentBlock {
|
||||
Text {
|
||||
@@ -148,7 +149,7 @@ pub enum CanonicalContentBlock {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalInstruction {
|
||||
pub role: CanonicalRole,
|
||||
#[serde(default)]
|
||||
@@ -157,7 +158,7 @@ pub struct CanonicalInstruction {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalMessage {
|
||||
pub role: CanonicalRole,
|
||||
#[serde(default)]
|
||||
@@ -166,7 +167,7 @@ pub struct CanonicalMessage {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalGenerationConfig {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u64>,
|
||||
@@ -192,7 +193,7 @@ pub struct CanonicalGenerationConfig {
|
||||
pub top_logprobs: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalToolDefinition {
|
||||
pub name: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -205,7 +206,7 @@ pub struct CanonicalToolDefinition {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalThinkingConfig {
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
@@ -215,7 +216,7 @@ pub struct CanonicalThinkingConfig {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalResponseFormat {
|
||||
pub format_type: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -224,7 +225,7 @@ pub struct CanonicalResponseFormat {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalUsage {
|
||||
#[serde(default)]
|
||||
pub input_tokens: u64,
|
||||
@@ -254,7 +255,7 @@ fn is_false(value: &bool) -> bool {
|
||||
!*value
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum CanonicalEmbeddingInput {
|
||||
String(String),
|
||||
@@ -264,7 +265,7 @@ pub enum CanonicalEmbeddingInput {
|
||||
Multimodal(Vec<CanonicalEmbeddingContent>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalEmbeddingContent {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<String>,
|
||||
@@ -336,7 +337,7 @@ impl CanonicalEmbeddingContent {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalEmbeddingRequest {
|
||||
pub input: CanonicalEmbeddingInput,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -353,7 +354,7 @@ pub struct CanonicalEmbeddingRequest {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalRerankRequest {
|
||||
pub query: String,
|
||||
#[serde(default)]
|
||||
@@ -374,7 +375,7 @@ impl CanonicalRerankRequest {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalEmbedding {
|
||||
#[serde(default)]
|
||||
pub index: usize,
|
||||
@@ -384,7 +385,7 @@ pub struct CanonicalEmbedding {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalEmbeddingResponse {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
@@ -396,7 +397,7 @@ pub struct CanonicalEmbeddingResponse {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalRequest {
|
||||
#[serde(default)]
|
||||
pub model: String,
|
||||
@@ -428,7 +429,7 @@ pub struct CanonicalRequest {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalResponseOutput {
|
||||
#[serde(default)]
|
||||
pub index: usize,
|
||||
@@ -454,7 +455,7 @@ impl Default for CanonicalResponseOutput {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CanonicalResponse {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
@@ -470,6 +471,425 @@ pub struct CanonicalResponse {
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
fn debug_json_bytes(value: &Value) -> Option<usize> {
|
||||
serde_json::to_vec(value).ok().map(|bytes| bytes.len())
|
||||
}
|
||||
|
||||
fn debug_json_map_bytes(value: &Map<String, Value>) -> Option<usize> {
|
||||
serde_json::to_vec(value).ok().map(|bytes| bytes.len())
|
||||
}
|
||||
|
||||
fn debug_json_option_bytes(value: Option<&Value>) -> Option<usize> {
|
||||
value.and_then(debug_json_bytes)
|
||||
}
|
||||
|
||||
fn debug_string_len(value: Option<&str>) -> Option<usize> {
|
||||
value.map(str::len)
|
||||
}
|
||||
|
||||
fn debug_string_list_summary(value: Option<&Vec<String>>) -> Option<(usize, usize)> {
|
||||
value.map(|values| (values.len(), values.iter().map(String::len).sum::<usize>()))
|
||||
}
|
||||
|
||||
fn debug_json_list_summary(value: &[Value]) -> (usize, usize) {
|
||||
(
|
||||
value.len(),
|
||||
value.iter().filter_map(debug_json_bytes).sum::<usize>(),
|
||||
)
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalToolChoice {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let mut debug = formatter.debug_struct("CanonicalToolChoice");
|
||||
match self {
|
||||
Self::Auto => debug.field("kind", &"auto"),
|
||||
Self::None => debug.field("kind", &"none"),
|
||||
Self::Required => debug.field("kind", &"required"),
|
||||
Self::Tool { name } => debug.field("kind", &"tool").field("name_len", &name.len()),
|
||||
}
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalContentBlock {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let mut debug = formatter.debug_struct("CanonicalContentBlock");
|
||||
match self {
|
||||
Self::Text { text, extensions } => debug
|
||||
.field("kind", &"text")
|
||||
.field("text_len", &text.len())
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::Thinking {
|
||||
text,
|
||||
signature,
|
||||
encrypted_content,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"thinking")
|
||||
.field("text_len", &text.len())
|
||||
.field("signature_len", &debug_string_len(signature.as_deref()))
|
||||
.field(
|
||||
"encrypted_content_len",
|
||||
&debug_string_len(encrypted_content.as_deref()),
|
||||
)
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::Image {
|
||||
data,
|
||||
url,
|
||||
media_type,
|
||||
detail,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"image")
|
||||
.field("data_len", &debug_string_len(data.as_deref()))
|
||||
.field("url_len", &debug_string_len(url.as_deref()))
|
||||
.field("media_type", media_type)
|
||||
.field("detail", detail)
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::File {
|
||||
data,
|
||||
file_id,
|
||||
file_url,
|
||||
media_type,
|
||||
filename,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"file")
|
||||
.field("data_len", &debug_string_len(data.as_deref()))
|
||||
.field("file_id_len", &debug_string_len(file_id.as_deref()))
|
||||
.field("file_url_len", &debug_string_len(file_url.as_deref()))
|
||||
.field("media_type", media_type)
|
||||
.field("filename_len", &debug_string_len(filename.as_deref()))
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::Audio {
|
||||
data,
|
||||
media_type,
|
||||
format,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"audio")
|
||||
.field("data_len", &debug_string_len(data.as_deref()))
|
||||
.field("media_type", media_type)
|
||||
.field("format", format)
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::ToolUse {
|
||||
id,
|
||||
name,
|
||||
input,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"tool_use")
|
||||
.field("id_len", &id.len())
|
||||
.field("name_len", &name.len())
|
||||
.field("input_bytes", &debug_json_bytes(input))
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::ToolResult {
|
||||
tool_use_id,
|
||||
name,
|
||||
output,
|
||||
content_text,
|
||||
is_error,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"tool_result")
|
||||
.field("tool_use_id_len", &tool_use_id.len())
|
||||
.field("name_len", &debug_string_len(name.as_deref()))
|
||||
.field("output_bytes", &debug_json_option_bytes(output.as_ref()))
|
||||
.field(
|
||||
"content_text_len",
|
||||
&debug_string_len(content_text.as_deref()),
|
||||
)
|
||||
.field("is_error", is_error)
|
||||
.field("extension_count", &extensions.len()),
|
||||
Self::Unknown {
|
||||
raw_type,
|
||||
payload,
|
||||
extensions,
|
||||
} => debug
|
||||
.field("kind", &"unknown")
|
||||
.field("raw_type_len", &raw_type.len())
|
||||
.field("payload_bytes", &debug_json_bytes(payload))
|
||||
.field("extension_count", &extensions.len()),
|
||||
}
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalInstruction {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalInstruction")
|
||||
.field("role", &self.role)
|
||||
.field("text_len", &self.text.len())
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalMessage {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalMessage")
|
||||
.field("role", &self.role)
|
||||
.field("content_count", &self.content.len())
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalGenerationConfig {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalGenerationConfig")
|
||||
.field("max_tokens", &self.max_tokens)
|
||||
.field("temperature", &self.temperature)
|
||||
.field("top_p", &self.top_p)
|
||||
.field("top_k", &self.top_k)
|
||||
.field(
|
||||
"stop_sequences",
|
||||
&debug_string_list_summary(self.stop_sequences.as_ref()),
|
||||
)
|
||||
.field("n", &self.n)
|
||||
.field("presence_penalty", &self.presence_penalty)
|
||||
.field("frequency_penalty", &self.frequency_penalty)
|
||||
.field("seed", &self.seed)
|
||||
.field("logprobs", &self.logprobs)
|
||||
.field("top_logprobs", &self.top_logprobs)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalToolDefinition {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalToolDefinition")
|
||||
.field("name_len", &self.name.len())
|
||||
.field(
|
||||
"description_len",
|
||||
&debug_string_len(self.description.as_deref()),
|
||||
)
|
||||
.field(
|
||||
"parameters_bytes",
|
||||
&debug_json_option_bytes(self.parameters.as_ref()),
|
||||
)
|
||||
.field("strict", &self.strict)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalThinkingConfig {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalThinkingConfig")
|
||||
.field("enabled", &self.enabled)
|
||||
.field("budget_tokens", &self.budget_tokens)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalResponseFormat {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalResponseFormat")
|
||||
.field("format_type_len", &self.format_type.len())
|
||||
.field(
|
||||
"json_schema_bytes",
|
||||
&debug_json_option_bytes(self.json_schema.as_ref()),
|
||||
)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalUsage {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalUsage")
|
||||
.field("input_tokens", &self.input_tokens)
|
||||
.field(
|
||||
"input_tokens_include_cache",
|
||||
&self.input_tokens_include_cache,
|
||||
)
|
||||
.field("output_tokens", &self.output_tokens)
|
||||
.field("total_tokens", &self.total_tokens)
|
||||
.field("cache_read_tokens", &self.cache_read_tokens)
|
||||
.field("cache_write_tokens", &self.cache_write_tokens)
|
||||
.field(
|
||||
"cache_creation_ephemeral_5m_tokens",
|
||||
&self.cache_creation_ephemeral_5m_tokens,
|
||||
)
|
||||
.field(
|
||||
"cache_creation_ephemeral_1h_tokens",
|
||||
&self.cache_creation_ephemeral_1h_tokens,
|
||||
)
|
||||
.field("reasoning_tokens", &self.reasoning_tokens)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalEmbeddingInput {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let mut debug = formatter.debug_struct("CanonicalEmbeddingInput");
|
||||
match self {
|
||||
Self::String(value) => debug
|
||||
.field("kind", &"string")
|
||||
.field("item_count", &1)
|
||||
.field("total_text_bytes", &value.len()),
|
||||
Self::StringArray(values) => debug
|
||||
.field("kind", &"string_array")
|
||||
.field("item_count", &values.len())
|
||||
.field(
|
||||
"total_text_bytes",
|
||||
&values.iter().map(String::len).sum::<usize>(),
|
||||
),
|
||||
Self::TokenArray(values) => debug
|
||||
.field("kind", &"token_array")
|
||||
.field("item_count", &values.len()),
|
||||
Self::TokenArrayArray(values) => debug
|
||||
.field("kind", &"token_array_array")
|
||||
.field("item_count", &values.len())
|
||||
.field(
|
||||
"total_token_count",
|
||||
&values.iter().map(Vec::len).sum::<usize>(),
|
||||
),
|
||||
Self::Multimodal(values) => debug
|
||||
.field("kind", &"multimodal")
|
||||
.field("item_count", &values.len()),
|
||||
}
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalEmbeddingContent {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalEmbeddingContent")
|
||||
.field("text_len", &debug_string_len(self.text.as_deref()))
|
||||
.field("image_len", &debug_string_len(self.image.as_deref()))
|
||||
.field("video_len", &debug_string_len(self.video.as_deref()))
|
||||
.field(
|
||||
"multi_images_summary",
|
||||
&self
|
||||
.multi_images
|
||||
.as_ref()
|
||||
.map(|images| (images.len(), images.iter().map(String::len).sum::<usize>())),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalEmbeddingRequest {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalEmbeddingRequest")
|
||||
.field("input", &self.input)
|
||||
.field("encoding_format", &self.encoding_format)
|
||||
.field("dimensions", &self.dimensions)
|
||||
.field("task_len", &debug_string_len(self.task.as_deref()))
|
||||
.field("user_len", &debug_string_len(self.user.as_deref()))
|
||||
.field(
|
||||
"parameters_bytes",
|
||||
&self.parameters.as_ref().and_then(debug_json_map_bytes),
|
||||
)
|
||||
.field("parameter_count", &self.parameters.as_ref().map(Map::len))
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalRerankRequest {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalRerankRequest")
|
||||
.field("query_len", &self.query.len())
|
||||
.field("documents", &debug_json_list_summary(&self.documents))
|
||||
.field("top_n", &self.top_n)
|
||||
.field("return_documents", &self.return_documents)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalEmbedding {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalEmbedding")
|
||||
.field("index", &self.index)
|
||||
.field("embedding_len", &self.embedding.len())
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalEmbeddingResponse {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalEmbeddingResponse")
|
||||
.field("id_len", &self.id.len())
|
||||
.field("model_len", &self.model.len())
|
||||
.field("embedding_count", &self.embeddings.len())
|
||||
.field("usage", &self.usage)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalRequest {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalRequest")
|
||||
.field("model_len", &self.model.len())
|
||||
.field("instruction_count", &self.instructions.len())
|
||||
.field("system_len", &debug_string_len(self.system.as_deref()))
|
||||
.field("message_count", &self.messages.len())
|
||||
.field("embedding", &self.embedding)
|
||||
.field("rerank", &self.rerank)
|
||||
.field("generation", &self.generation)
|
||||
.field("tool_count", &self.tools.len())
|
||||
.field("tool_choice", &self.tool_choice)
|
||||
.field("thinking", &self.thinking)
|
||||
.field("response_format", &self.response_format)
|
||||
.field("parallel_tool_calls", &self.parallel_tool_calls)
|
||||
.field(
|
||||
"metadata_bytes",
|
||||
&debug_json_option_bytes(self.metadata.as_ref()),
|
||||
)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalResponseOutput {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalResponseOutput")
|
||||
.field("index", &self.index)
|
||||
.field("role", &self.role)
|
||||
.field("content_count", &self.content.len())
|
||||
.field("stop_reason", &self.stop_reason)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CanonicalResponse {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CanonicalResponse")
|
||||
.field("id_len", &self.id.len())
|
||||
.field("model_len", &self.model.len())
|
||||
.field("output_count", &self.outputs.len())
|
||||
.field("content_count", &self.content.len())
|
||||
.field("stop_reason", &self.stop_reason)
|
||||
.field("usage", &self.usage)
|
||||
.field("extension_count", &self.extensions.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_openai_chat_to_canonical_request(body_json: &Value) -> Option<CanonicalRequest> {
|
||||
crate::formats::openai::chat::request::from_raw(body_json)
|
||||
}
|
||||
@@ -4531,7 +4951,7 @@ pub(crate) type GeminiCanonicalTools = (
|
||||
Option<Value>,
|
||||
);
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct GeminiGoogleSearchGrounding {
|
||||
pub source_field: &'static str,
|
||||
pub source_dialect: &'static str,
|
||||
@@ -4541,6 +4961,23 @@ pub(crate) struct GeminiGoogleSearchGrounding {
|
||||
pub output_payload: Value,
|
||||
}
|
||||
|
||||
impl fmt::Debug for GeminiGoogleSearchGrounding {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("GeminiGoogleSearchGrounding")
|
||||
.field("source_field", &self.source_field)
|
||||
.field("source_dialect", &self.source_dialect)
|
||||
.field("legacy", &self.legacy)
|
||||
.field("payload_bytes", &debug_json_bytes(&self.payload))
|
||||
.field("raw_payload_bytes", &debug_json_bytes(&self.raw_payload))
|
||||
.field(
|
||||
"output_payload_bytes",
|
||||
&debug_json_bytes(&self.output_payload),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn gemini_google_search_grounding(
|
||||
tool_object: &Map<String, Value>,
|
||||
) -> Option<GeminiGoogleSearchGrounding> {
|
||||
|
||||
@@ -256,6 +256,10 @@ fn transform_provider_private_stream_line_with_event_state(
|
||||
|
||||
const CONNECT_FRAME_HEADER_BYTES: usize = 5;
|
||||
const MAX_CONNECT_JSON_FRAME_BYTES: usize = 16 * 1024 * 1024;
|
||||
// Keep malformed or incomplete provider streams from growing this parser's
|
||||
// carry buffer without bound when no complete SSE/Connect record arrives.
|
||||
const MAX_PRIVATE_STREAM_BUFFER_BYTES: usize =
|
||||
MAX_CONNECT_JSON_FRAME_BYTES + CONNECT_FRAME_HEADER_BYTES;
|
||||
|
||||
fn report_context_is_windsurf_envelope(report_context: &Value) -> bool {
|
||||
report_context
|
||||
@@ -424,6 +428,18 @@ impl ProviderPrivateStreamNormalizer<'_> {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
|
||||
let next_len = self
|
||||
.buffered
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.ok_or_else(|| {
|
||||
AiSurfaceFinalizeError::new("provider stream normalization buffer overflow")
|
||||
})?;
|
||||
if next_len > MAX_PRIVATE_STREAM_BUFFER_BYTES {
|
||||
return Err(AiSurfaceFinalizeError::new(format!(
|
||||
"provider stream normalization buffer exceeds {MAX_PRIVATE_STREAM_BUFFER_BYTES} bytes"
|
||||
)));
|
||||
}
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
if report_context_is_windsurf_envelope(self.report_context)
|
||||
&& buffer_looks_like_connect_frame(&self.buffered)
|
||||
@@ -1282,4 +1298,21 @@ data: {"type":"response.failed","response":{"status":"failed","error":{"message"
|
||||
assert!(output_text.contains("\"message\":\"rate limited\""));
|
||||
assert!(!output_text.contains("chat.completion.chunk"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn private_stream_normalizer_rejects_unbounded_incomplete_buffer() {
|
||||
let report_context = json!({
|
||||
"has_envelope": true,
|
||||
"envelope_name": "antigravity:v1internal",
|
||||
"provider_api_format": "gemini:generate_content",
|
||||
});
|
||||
let mut normalizer = maybe_build_provider_private_stream_normalizer(Some(&report_context))
|
||||
.expect("normalizer should exist");
|
||||
let oversized = vec![b'x'; super::MAX_PRIVATE_STREAM_BUFFER_BYTES + 1];
|
||||
|
||||
let error = normalizer
|
||||
.push_chunk(&oversized)
|
||||
.expect_err("incomplete provider stream must be bounded");
|
||||
assert!(error.0.contains("buffer exceeds"));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user