Browse Source

fix: return Failure envelope when SSRF blocks AI chat URL at runtime

Wrap configure() call in AiChatModelServiceImpl.sendChatRequestAsync
so a synchronous exception (e.g. runtime SSRF block) is converted to
a failed future and caught by the controller's .catching() chain,
returning TbChatResponse.Failure instead of a raw 500.
pull/15412/head
Oleksandra Matviienko 6 months ago
parent
commit
d8c1a91721
  1. 8
      application/src/main/java/org/thingsboard/server/service/ai/AiChatModelServiceImpl.java
  2. 37
      application/src/test/java/org/thingsboard/server/controller/AiModelControllerTest.java

8
application/src/main/java/org/thingsboard/server/service/ai/AiChatModelServiceImpl.java

@ -17,6 +17,7 @@ package org.thingsboard.server.service.ai;
import com.fasterxml.jackson.core.io.JsonStringEncoder;
import com.google.common.util.concurrent.FluentFuture;
import com.google.common.util.concurrent.Futures;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.TextContent;
@ -42,7 +43,12 @@ class AiChatModelServiceImpl implements AiChatModelService {
@Override
public <C extends AiChatModelConfig<C>> FluentFuture<ChatResponse> sendChatRequestAsync(AiChatModelConfig<C> chatModelConfig, ChatRequest chatRequest) {
ChatModel langChainChatModel = chatModelConfig.configure(chatModelConfigurer);
ChatModel langChainChatModel;
try {
langChainChatModel = chatModelConfig.configure(chatModelConfigurer);
} catch (Throwable t) {
return FluentFuture.from(Futures.immediateFailedFuture(t));
}
if (langChainChatModel.provider() == ModelProvider.GITHUB_MODELS) {
chatRequest = prepareGithubChatRequest(chatRequest);
}

37
application/src/test/java/org/thingsboard/server/controller/AiModelControllerTest.java

@ -22,6 +22,10 @@ import org.springframework.test.web.servlet.ResultActions;
import org.thingsboard.common.util.SsrfProtectionValidator;
import org.thingsboard.server.common.data.EntityType;
import org.thingsboard.server.common.data.ai.AiModel;
import org.thingsboard.server.common.data.ai.dto.TbChatRequest;
import org.thingsboard.server.common.data.ai.dto.TbChatResponse;
import org.thingsboard.server.common.data.ai.dto.TbContent;
import org.thingsboard.server.common.data.ai.dto.TbUserMessage;
import org.thingsboard.server.common.data.ai.model.chat.AnthropicChatModelConfig;
import org.thingsboard.server.common.data.ai.model.chat.GoogleAiGeminiChatModelConfig;
import org.thingsboard.server.common.data.ai.model.chat.OpenAiChatModelConfig;
@ -35,6 +39,8 @@ import org.thingsboard.server.common.data.page.PageLink;
import org.thingsboard.server.common.data.page.SortOrder;
import org.thingsboard.server.dao.service.DaoSqlTest;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
@ -168,6 +174,37 @@ public class AiModelControllerTest extends AbstractControllerTest {
}
}
@Test
public void sendChatRequest_whenBaseUrlBlockedAtRuntime_shouldReturnFailureEnvelope() throws Exception {
// GIVEN
loginTenantAdmin();
SsrfProtectionValidator.setEnabled(true);
try {
var modelConfig = OpenAiChatModelConfig.builder()
.providerConfig(OpenAiProviderConfig.builder()
.baseUrl("http://10.0.0.1:8080/")
.apiKey("test-api-key")
.build())
.modelId("gpt-4o")
.build();
var chatRequest = new TbChatRequest(
null,
new TbUserMessage(List.of(new TbContent.TbTextContent("hi"))),
modelConfig);
// WHEN
TbChatResponse response = doPost("/api/ai/chat", chatRequest, TbChatResponse.class);
// THEN
assertThat(response).isInstanceOf(TbChatResponse.Failure.class);
assertThat(((TbChatResponse.Failure) response).errorDetails()).contains("URI is invalid");
} finally {
SsrfProtectionValidator.setEnabled(false);
}
}
/* --- Get by ID API tests --- */
@Test

Loading…
Cancel
Save