diff --git a/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelCredentialPolicy.java b/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelCredentialPolicy.java index 1409e9f6..4a462cd9 100644 --- a/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelCredentialPolicy.java +++ b/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelCredentialPolicy.java @@ -16,6 +16,12 @@ public final class ChatModelCredentialPolicy { public static final String DEEPSEEK_API_HOST = "https://api.deepseek.com"; public static final String PPIO_PROVIDER = "ppio"; public static final String PPIO_API_HOST = "https://api.ppio.com/openai/v1"; + public static final String ATLAS_PROVIDER = "atlas"; + private static final Set ATLAS_API_HOSTS = Set.of( + "https://api.atlascloud.ai", "https://api.atlascloud.ai/", + "https://api.atlascloud.ai/v1", "https://api.atlascloud.ai/v1/", + "https://api.atlascloud.ai/api/v1", "https://api.atlascloud.ai/api/v1/" + ); private static final Set PPIO_API_HOSTS = Set.of( PPIO_API_HOST, PPIO_API_HOST + "/", @@ -69,6 +75,10 @@ public final class ChatModelCredentialPolicy { String apiHost, String apiKey) { requireSecureApiHost(apiHost); + if (ATLAS_PROVIDER.equals(providerCode)) { + requireAtlasConfiguration(providerCode, modelName, apiHost, apiKey); + return; + } if (CustomApiCredentialPolicy.isCustomProvider(providerCode)) { CustomApiCredentialPolicy.requireConfiguration(providerCode, modelName, apiHost, apiKey); return; @@ -121,6 +131,8 @@ public final class ChatModelCredentialPolicy { expectedReference = "env:DEEPSEEK_API_KEY"; } else if (PPIO_PROVIDER.equals(providerCode)) { expectedReference = "env:PPIO_API_KEY"; + } else if (ATLAS_PROVIDER.equals(providerCode)) { + expectedReference = "env:ATLAS_API_KEY"; } else { throw new IllegalArgumentException("Model credential provider is not allowlisted"); } @@ -137,11 +149,29 @@ public final class ChatModelCredentialPolicy { CustomApiCredentialPolicy.requireConfiguration(providerCode, modelName, apiHost, apiKey); } else if (PPIO_PROVIDER.equals(providerCode)) { requirePpioConfiguration(providerCode, modelName, apiHost, apiKey); + } else if (ATLAS_PROVIDER.equals(providerCode)) { + requireAtlasConfiguration(providerCode, modelName, apiHost, apiKey); } else { requireDeepSeekConfiguration(providerCode, modelName, apiHost, apiKey); } } + /** Atlas credentials may only be sent to its official media API origin. */ + public static void requireAtlasConfiguration(String providerCode, String modelName, + String apiHost, String apiKey) { + if (!ATLAS_PROVIDER.equals(providerCode)) { + throw new IllegalArgumentException("Atlas model provider is not trusted"); + } + if (modelName == null || !modelName.matches("[A-Za-z0-9][A-Za-z0-9._/-]{0,254}")) { + throw new IllegalArgumentException("Atlas model ID is invalid"); + } + requireSecureApiHost(apiHost); + if (!ATLAS_API_HOSTS.contains(apiHost)) { + throw new IllegalArgumentException("Atlas API host is not allowlisted"); + } + requireProviderReference(providerCode, apiKey); + } + /** Validates PPIO and returns the base URL expected by the OpenAI Chat Completions client. */ public static String requirePpioConfiguration(String providerCode, String modelName, String apiHost, String apiKey) { diff --git a/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelSecretReference.java b/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelSecretReference.java index 40a08ccb..ba7b16ef 100644 --- a/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelSecretReference.java +++ b/ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelSecretReference.java @@ -8,7 +8,7 @@ import java.util.regex.Pattern; public final class ChatModelSecretReference { public static final String ENV_REFERENCE_REGEXP = - "^env:(?:(?:DEEPSEEK|PPIO)_API_KEY|" + CustomApiCredentialPolicy.REFERENCE_PATTERN + ")$"; + "^env:(?:(?:DEEPSEEK|PPIO|ATLAS)_API_KEY|" + CustomApiCredentialPolicy.REFERENCE_PATTERN + ")$"; private static final String ENV_PREFIX = "env:"; private static final Pattern ENV_REFERENCE = Pattern.compile(ENV_REFERENCE_REGEXP); diff --git a/ruoyi-common/ruoyi-common-chat/src/test/java/org/ruoyi/common/chat/security/AtlasCredentialPolicyTest.java b/ruoyi-common/ruoyi-common-chat/src/test/java/org/ruoyi/common/chat/security/AtlasCredentialPolicyTest.java new file mode 100644 index 00000000..561fb7c7 --- /dev/null +++ b/ruoyi-common/ruoyi-common-chat/src/test/java/org/ruoyi/common/chat/security/AtlasCredentialPolicyTest.java @@ -0,0 +1,50 @@ +package org.ruoyi.common.chat.security; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.*; + +@Tag("dev") +class AtlasCredentialPolicyTest { + private static final String MODEL = "bytedance/seed-audio-1.0"; + private static final String HOST = "https://api.atlascloud.ai/v1"; + private static final String REF = "env:ATLAS_API_KEY"; + + @Test + void acceptsOfficialEndpointFormsAtSaveAndUse() { + for (String suffix : new String[]{"", "/", "/v1", "/v1/", "/api/v1", "/api/v1/"}) { + String host = "https://api.atlascloud.ai" + suffix; + assertDoesNotThrow(() -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, host, REF)); + assertEquals("test-secret", ChatModelCredentialPolicy.resolveApiKeyForUse("atlas", "atlas", MODEL, host, REF, + name -> "ATLAS_API_KEY".equals(name) ? "test-secret" : null)); + } + } + + @Test + void rejectsUntrustedEndpointsBeforeReadingEnvironment() { + for (String host : new String[]{"http://api.atlascloud.ai/v1", "https://evil.example/v1", + "https://api.atlascloud.ai.evil.example/v1", "https://api.atlascloud.ai:8443/v1", + "https://user@api.atlascloud.ai/v1", "https://api.atlascloud.ai/v1?x=1", "https://api.atlascloud.ai/other"}) { + assertThrows(IllegalArgumentException.class, + () -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, host, REF)); + assertThrows(IllegalArgumentException.class, + () -> ChatModelCredentialPolicy.resolveApiKeyForUse("atlas", "atlas", MODEL, host, REF, + name -> { fail("Must validate endpoint before resolving secret"); return null; })); + } + } + + @Test + void rejectsWrongProviderReferenceAndMissingSecret() { + assertThrows(IllegalArgumentException.class, () -> ChatModelCredentialPolicy.resolveApiKeyForUse( + "openai", "atlas", MODEL, HOST, REF, name -> "secret")); + for (String key : new String[]{"env:PPIO_API_KEY", "plaintext", "", null}) { + assertThrows(IllegalArgumentException.class, + () -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, HOST, key)); + } + assertThrows(IllegalStateException.class, () -> ChatModelCredentialPolicy.resolveApiKeyForUse( + "atlas", "atlas", MODEL, HOST, REF, name -> null)); + assertThrows(IllegalArgumentException.class, + () -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", "bad model", HOST, REF)); + } +} diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaContentController.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaContentController.java new file mode 100644 index 00000000..6c540e92 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaContentController.java @@ -0,0 +1,44 @@ +package org.ruoyi.controller.chat; + +import jakarta.validation.constraints.NotBlank; +import lombok.RequiredArgsConstructor; +import org.ruoyi.common.chat.service.chat.IChatModelService; +import org.ruoyi.service.media.AtlasMediaContentService; +import org.springframework.http.CacheControl; +import org.springframework.http.HttpHeaders; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.validation.annotation.Validated; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RequestParam; +import org.springframework.web.bind.annotation.RestController; + +import java.io.IOException; + +@Validated +@RestController +@RequiredArgsConstructor +@RequestMapping("/media") +public class MediaContentController { + private final IChatModelService models; + private final AtlasMediaContentService content; + + /** Uses the normal login interceptor; browser clients fetch with their Authorization header. */ + @GetMapping("/content") + public ResponseEntity content(@RequestParam @NotBlank String model, + @RequestParam @NotBlank String predictionId) throws IOException { + var configuredModel = models.selectModelByName(model); + if (configuredModel == null) { + throw new IllegalArgumentException("未找到模型配置"); + } + var result = content.retrieve(configuredModel, predictionId); + return ResponseEntity.ok() + .contentType(MediaType.parseMediaType(result.mimeType())) + .contentLength(result.bytes().length) + .cacheControl(CacheControl.noStore()) + .header(HttpHeaders.CONTENT_DISPOSITION, "inline") + .header("X-Content-Type-Options", "nosniff") + .body(result.bytes()); + } +} diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaGenerationController.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaGenerationController.java index b2f032ab..5bf7ea08 100644 --- a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaGenerationController.java +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaGenerationController.java @@ -56,14 +56,18 @@ public class MediaGenerationController { @PostMapping("/image") public R image(@Valid @RequestBody ImageGenerationRequest request) { ChatModelVo model = loadModel(request.getModel(), ModelType.IMAGE.getKey()); - String result = imageServiceFactory.getOriginalService(model.getProviderCode()) - .generateImage(ImageContext.builder() - .chatModelVo(model) - .prompt(request.getPrompt()) - .size(request.getSize()) - .seed(request.getSeed()) - .build()); - return R.ok(toImageResponse(result)); + ImageContext context = ImageContext.builder() + .chatModelVo(model) + .prompt(request.getPrompt()) + .size(request.getSize()) + .seed(request.getSeed()) + .build(); + var service = imageServiceFactory.getOriginalService(model.getProviderCode()); + // Atlas sync requests can outlive its gateway timeout. Return a task for the workbench to poll. + if ("atlas".equals(model.getProviderCode())) { + return R.ok(service.startImageGeneration(context)); + } + return R.ok(toImageResponse(service.generateImage(context))); } @PostMapping("/video") diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasMediaContentService.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasMediaContentService.java new file mode 100644 index 00000000..75b7ecd5 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasMediaContentService.java @@ -0,0 +1,97 @@ +package org.ruoyi.service.media; + +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Service; + +import java.io.IOException; +import java.util.Locale; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +/** Delivers Atlas task outputs without exposing a general-purpose URL proxy. */ +@Service +public class AtlasMediaContentService { + static final int MAX_BYTES = 64 * 1024 * 1024; + private static final Set OUTPUT_HOSTS = Set.of( + "atlas-media.oss-us-west-1.aliyuncs.com", + "ark-acg-ap-southeast-1.tos-ap-southeast-1.volces.com" + ); + private static final Set MIME_TYPES = Set.of( + "image/jpeg", "image/png", "image/webp", "image/gif", + "audio/mpeg", "audio/wav", "audio/x-wav", "audio/ogg", "audio/mp4", "audio/flac", + "video/mp4", "video/webm" + ); + + private final AtlasPredictionService predictions; + private final OkHttpClient client; + + @Autowired + public AtlasMediaContentService(AtlasPredictionService predictions) { + this(predictions, new OkHttpClient.Builder() + .followRedirects(false).followSslRedirects(false) + .connectTimeout(15, TimeUnit.SECONDS) + .readTimeout(60, TimeUnit.SECONDS) + .callTimeout(90, TimeUnit.SECONDS) + .build()); + } + + AtlasMediaContentService(AtlasPredictionService predictions, OkHttpClient client) { + this.predictions = predictions; + this.client = client; + } + + public Content retrieve(ChatModelVo model, String predictionId) throws IOException { + if (!"atlas".equals(model.getProviderCode()) + || !Set.of("image", "audio", "video").contains(model.getCategory())) { + throw new IllegalArgumentException("此模型不支持 Atlas 资源预览"); + } + if (predictionId == null || !predictionId.matches("[A-Za-z0-9_-]{1,128}")) { + throw new IllegalArgumentException("无效的任务 ID"); + } + // Resolve the output using this configured model's credentials, never a caller-supplied URL. + var result = predictions.retrieve(model, predictionId); + if (!Set.of("completed", "succeeded", "success").contains(String.valueOf(result.getStatus())) + || !predictionId.equals(result.getId())) { + throw new IllegalArgumentException("任务尚未完成、已过期或不存在"); + } + HttpUrl url = checkedOutputUrl(result.getUrl()); + // The API key belongs only on the prediction API request, never on the media CDN request. + Request request = new Request.Builder().url(url).get().build(); + try (Response response = client.newCall(request).execute()) { + ResponseBody body = response.body(); + if (response.code() != 200 || body == null) { + throw new IOException("资源服务暂不可用，请稍后重试"); + } + String mime = response.header("Content-Type", "").split(";", 2)[0].trim().toLowerCase(Locale.ROOT); + if (!MIME_TYPES.contains(mime) || !mime.startsWith(model.getCategory() + "/")) { + throw new IOException("资源格式与模型类型不匹配或不支持预览"); + } + if (body.contentLength() > MAX_BYTES) { + throw new IOException("资源超过 64 MiB 预览上限"); + } + byte[] bytes = body.byteStream().readNBytes(MAX_BYTES + 1); + if (bytes.length == 0 || bytes.length > MAX_BYTES) { + throw new IOException("资源为空或超过 64 MiB 预览上限"); + } + return new Content(bytes, mime); + } + } + + static HttpUrl checkedOutputUrl(String value) { + HttpUrl url = value == null ? null : HttpUrl.parse(value); + if (url == null || !url.isHttps() || url.port() != 443 + || !url.username().isEmpty() || !url.password().isEmpty() || url.fragment() != null + || !OUTPUT_HOSTS.contains(url.host())) { + throw new IllegalArgumentException("资源地址不在 Atlas 预览白名单内"); + } + return url; + } + + public record Content(byte[] bytes, String mimeType) { } +} diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasPredictionService.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasPredictionService.java index 6e8f12a9..d688979a 100644 --- a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasPredictionService.java +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasPredictionService.java @@ -43,7 +43,8 @@ public class AtlasPredictionService { log.warn("Atlas Cloud 预测任务不存在或已过期: predictionId={}, category={}", predictionId, model.getCategory()); return MediaGenerationResponse.builder() .type(mediaType(model.getCategory())) - .mimeType("image".equals(model.getCategory()) ? "image/png" : "video/mp4") + .mimeType(mimeType(mediaType(model.getCategory()), null)) + .id(predictionId) .status("failed") .rawResponse(responseText) .build(); @@ -71,7 +72,7 @@ public class AtlasPredictionService { } return MediaGenerationResponse.builder() .type(type) - .mimeType("image".equals(type) ? "image/png" : "video/mp4") + .mimeType(mimeType(type, firstOutput(data))) .id(AtlasMediaSupport.text(data, "id")) .status(status) .url(firstOutput(data)) @@ -80,6 +81,22 @@ public class AtlasPredictionService { .build(); } + private String mimeType(String type, String url) { + String path = url == null ? "" : url.split("[?#]", 2)[0].toLowerCase(java.util.Locale.ROOT); + if ("audio".equals(type)) { + if (path.endsWith(".wav")) return "audio/wav"; + if (path.endsWith(".ogg") || path.endsWith(".opus")) return "audio/ogg"; + if (path.endsWith(".pcm")) return "application/octet-stream"; + return "audio/mpeg"; + } + if ("image".equals(type)) { + if (path.endsWith(".jpg") || path.endsWith(".jpeg")) return "image/jpeg"; + if (path.endsWith(".webp")) return "image/webp"; + return "image/png"; + } + return "video/mp4"; + } + /** * 提取末帧 URL。Atlas return_last_frame=true 时会在 outputs 或顶层节点返回末帧图片， * 字段名兼容 last_frame_url / end_frame_url / last_frame / last_frame_image。 @@ -128,7 +145,11 @@ public class AtlasPredictionService { return null; } - private String mediaType(String category) { - return "image".equals(category) ? "image" : "video"; + static String mediaType(String category) { + return switch (category) { + case "image" -> "image"; + case "audio" -> "audio"; + default -> "video"; + }; } } diff --git a/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaContentControllerTest.java b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaContentControllerTest.java new file mode 100644 index 00000000..081e332f --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaContentControllerTest.java @@ -0,0 +1,29 @@ +package org.ruoyi.controller.chat; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo; +import org.ruoyi.common.chat.service.chat.IChatModelService; +import org.ruoyi.service.media.AtlasMediaContentService; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +@Tag("dev") +class MediaContentControllerTest { + @Test + void returnsInlineNonCachedBytes() throws Exception { + var models = mock(IChatModelService.class); + var service = mock(AtlasMediaContentService.class); + var model = new ChatModelVo(); + when(models.selectModelByName("image-model")).thenReturn(model); + when(service.retrieve(model, "task-123")).thenReturn(new AtlasMediaContentService.Content(new byte[]{1, 2}, "image/jpeg")); + var result = new MediaContentController(models, service).content("image-model", "task-123"); + assertEquals("image/jpeg", result.getHeaders().getContentType().toString()); + assertEquals("inline", result.getHeaders().getFirst("Content-Disposition")); + assertEquals("no-store", result.getHeaders().getCacheControl()); + assertEquals("nosniff", result.getHeaders().getFirst("X-Content-Type-Options")); + assertEquals(2, result.getHeaders().getContentLength()); + assertArrayEquals(new byte[]{1, 2}, result.getBody()); + } +} diff --git a/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaGenerationControllerTest.java b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaGenerationControllerTest.java new file mode 100644 index 00000000..45804425 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaGenerationControllerTest.java @@ -0,0 +1,65 @@ +package org.ruoyi.controller.chat; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo; +import org.ruoyi.common.chat.entity.media.MediaGenerationResponse; +import org.ruoyi.common.chat.factory.ImageServiceFactory; +import org.ruoyi.common.chat.service.chat.IChatModelService; +import org.ruoyi.common.chat.service.image.IImageGenerationService; +import org.ruoyi.domain.bo.media.ImageGenerationRequest; +import org.ruoyi.service.media.AtlasPredictionService; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +@Tag("dev") +class MediaGenerationControllerTest { + @Test + void returnsAtlasTaskWithoutBlockingForImage() { + var models = mock(IChatModelService.class); + var factory = mock(ImageServiceFactory.class); + var service = mock(IImageGenerationService.class); + var model = model("atlas"); + when(models.selectModelByName("test-image")).thenReturn(model); + when(factory.getOriginalService("atlas")).thenReturn(service); + when(service.startImageGeneration(any())).thenReturn(MediaGenerationResponse.builder() + .id("task-123").type("image").status("pending").build()); + var controller = new MediaGenerationController(models, null, factory, null, new AtlasPredictionService()); + var result = controller.image(request()).getData(); + assertEquals("task-123", result.getId()); + assertEquals("pending", result.getStatus()); + verify(service, never()).generateImage(any()); + } + + @Test + void preservesOtherProvidersBase64ImageResponse() { + var models = mock(IChatModelService.class); + var factory = mock(ImageServiceFactory.class); + var service = mock(IImageGenerationService.class); + when(models.selectModelByName("test-image")).thenReturn(model("openai")); + when(factory.getOriginalService("openai")).thenReturn(service); + when(service.generateImage(any())).thenReturn("data:image/png;base64,aGVsbG8="); + var controller = new MediaGenerationController(models, null, factory, null, new AtlasPredictionService()); + var result = controller.image(request()).getData(); + assertEquals("image/png", result.getMimeType()); + assertEquals("aGVsbG8=", result.getB64Json()); + verify(service, never()).startImageGeneration(any()); + } + + private ChatModelVo model(String provider) { + var model = new ChatModelVo(); + model.setModelName("test-image"); + model.setProviderCode(provider); + model.setCategory("image"); + return model; + } + + private ImageGenerationRequest request() { + var request = new ImageGenerationRequest(); + request.setModel("test-image"); + request.setPrompt("a blue book"); + return request; + } +} diff --git a/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasMediaContentServiceTest.java b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasMediaContentServiceTest.java new file mode 100644 index 00000000..1d10b884 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasMediaContentServiceTest.java @@ -0,0 +1,119 @@ +package org.ruoyi.service.media; + +import okhttp3.Call; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo; +import org.ruoyi.common.chat.entity.media.MediaGenerationResponse; + +import java.io.IOException; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +@Tag("dev") +class AtlasMediaContentServiceTest { + private static final String URL = "https://atlas-media.oss-us-west-1.aliyuncs.com/assetd-history/result.mp3"; + private final AtlasPredictionService predictions = mock(AtlasPredictionService.class); + private final OkHttpClient client = mock(OkHttpClient.class); + private final AtlasMediaContentService service = new AtlasMediaContentService(predictions, client); + + @Test + void servesMediaBytesWithoutForwardingCredentialsOrDownloadHeaders() throws Exception { + var model = model(); + model.setApiKey("env:ATLAS_API_KEY"); + completed(model); + respond(200, ResponseBody.create("ID3audio", MediaType.get("audio/mpeg")), "audio/mpeg"); + var content = service.retrieve(model, "task-123"); + assertEquals("audio/mpeg", content.mimeType()); + assertArrayEquals("ID3audio".getBytes(), content.bytes()); + var capture = ArgumentCaptor.forClass(Request.class); + verify(client).newCall(capture.capture()); + assertEquals(URL, capture.getValue().url().toString()); + assertNull(capture.getValue().header("Authorization")); + assertNull(capture.getValue().header("Cookie")); + } + + @Test + void rejectsUntrustedUrlsAndNonHttpsVariants() { + for (String url : new String[]{"http://atlas-media.oss-us-west-1.aliyuncs.com/a.mp3", + "https://127.0.0.1/a", "https://atlas-media.oss-us-west-1.aliyuncs.com.evil.test/a", + "https://user@atlas-media.oss-us-west-1.aliyuncs.com/a", + "https://atlas-media.oss-us-west-1.aliyuncs.com:8443/a", + "https://atlas-media.oss-us-west-1.aliyuncs.com/a#fragment", "file:///etc/passwd"}) { + assertThrows(IllegalArgumentException.class, () -> AtlasMediaContentService.checkedOutputUrl(url)); + } + assertEquals("atlas-media.oss-us-west-1.aliyuncs.com", AtlasMediaContentService.checkedOutputUrl(URL).host()); + } + + @Test + void rejectsInvalidIdsAndOtherProvidersBeforeMakingRequests() { + var model = model(); + assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "../task?url=bad")); + model.setProviderCode("openai"); + assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123")); + verifyNoInteractions(predictions, client); + } + + @Test + void rejectsUnfinishedOrMismatchedTasksBeforeFetchingContent() { + var model = model(); + when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder() + .id("task-123").status("processing").url(URL).build()); + assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123")); + when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder() + .id("different").status("completed").url(URL).build()); + assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123")); + verifyNoInteractions(client); + } + + @Test + void rejectsRedirectsAndHtmlResponses() throws Exception { + var model = model(); + completed(model); + respond(302, ResponseBody.create("", MediaType.get("audio/mpeg")), "audio/mpeg"); + assertThrows(IOException.class, () -> service.retrieve(model, "task-123")); + respond(200, ResponseBody.create("error", MediaType.get("text/html")), "text/html"); + assertThrows(IOException.class, () -> service.retrieve(model, "task-123")); + } + + @Test + void rejectsOversizedContentBeforeReadingTheBody() throws Exception { + var model = model(); + completed(model); + var body = mock(ResponseBody.class); + when(body.contentLength()).thenReturn((long) AtlasMediaContentService.MAX_BYTES + 1); + respond(200, body, "audio/mpeg"); + assertThrows(IOException.class, () -> service.retrieve(model, "task-123")); + verify(body, never()).byteStream(); + } + + private ChatModelVo model() { + var model = new ChatModelVo(); + model.setCategory("audio"); + model.setProviderCode("atlas"); + return model; + } + + private void completed(ChatModelVo model) { + when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder() + .id("task-123").status("completed").url(URL).build()); + } + + private void respond(int status, ResponseBody body, String type) throws IOException { + Call call = mock(Call.class); + when(client.newCall(any())).thenReturn(call); + when(call.execute()).thenReturn(new Response.Builder().request(new Request.Builder().url(URL).build()) + .protocol(Protocol.HTTP_1_1).code(status).message("test").body(body) + .header("Content-Type", type).header("Content-Disposition", "attachment") + .header("x-oss-force-download", "true").header("Location", "http://127.0.0.1/private").build()); + } +} diff --git a/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasPredictionServiceTest.java b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasPredictionServiceTest.java new file mode 100644 index 00000000..6f6b93b9 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasPredictionServiceTest.java @@ -0,0 +1,26 @@ +package org.ruoyi.service.media; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Tag("dev") +class AtlasPredictionServiceTest { + @Test + void keepsAudioTypeAndMimeThroughPolling() throws Exception { + var service = new AtlasPredictionService(); + assertEquals("audio", AtlasPredictionService.mediaType("audio")); + assertEquals("image", AtlasPredictionService.mediaType("image")); + assertEquals("video", AtlasPredictionService.mediaType("video")); + var pending = service.toResponse("{\"data\":{\"id\":\"task-123\",\"status\":\"processing\"}}", "audio"); + assertEquals("audio", pending.getType()); + assertEquals("audio/mpeg", pending.getMimeType()); + assertEquals("task-123", pending.getId()); + for (String[] sample : new String[][]{{"mp3", "audio/mpeg"}, {"wav", "audio/wav"}, {"ogg", "audio/ogg"}}) { + var complete = service.toResponse("{\"data\":{\"status\":\"completed\",\"outputs\":[\"https://example.com/result." + + sample[0] + "?token=test\"]}}", "audio"); + assertEquals(sample[1], complete.getMimeType()); + } + } +}