From 10c2e53c27e37343c376e817cb8da7d74945c355 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 26 Aug 2026 10:38:56 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=AE=20fix:=20Screen=20User-Supplied=20?= =?UTF-8?q?Structured=20Tool=20Endpoints=20(#15253)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../clients/tools/structured/AzureAISearch.js | 23 ++++++ .../tools/structured/AzureAISearch.spec.js | 54 ++++++++++++++ .../tools/structured/StableDiffusion.js | 9 ++- .../tools/structured/StableDiffusion.spec.js | 70 +++++++++++++++++++ api/app/clients/tools/util/handleTools.js | 16 ++++- .../clients/tools/util/handleTools.test.js | 15 ++++ 6 files changed, 184 insertions(+), 3 deletions(-) create mode 100644 api/app/clients/tools/structured/AzureAISearch.spec.js create mode 100644 api/app/clients/tools/structured/StableDiffusion.spec.js diff --git a/api/app/clients/tools/structured/AzureAISearch.js b/api/app/clients/tools/structured/AzureAISearch.js index 0579614341..0a74245e8e 100644 --- a/api/app/clients/tools/structured/AzureAISearch.js +++ b/api/app/clients/tools/structured/AzureAISearch.js @@ -2,6 +2,21 @@ const { logger } = require('@librechat/data-schemas'); const { Tool } = require('@librechat/agents/langchain/tools'); const { SearchClient, AzureKeyCredential } = require('@azure/search-documents'); +const azureSearchHostSuffixes = ['.search.windows.net', '.search.azure.us', '.search.azure.cn']; + +const isAzureSearchEndpoint = (endpoint) => { + try { + const parsed = new URL(endpoint); + const hostname = parsed.hostname.toLowerCase().replace(/\.$/, ''); + return ( + parsed.protocol === 'https:' && + azureSearchHostSuffixes.some((suffix) => hostname.endsWith(suffix)) + ); + } catch { + return false; + } +}; + const azureAISearchJsonSchema = { type: 'object', properties: { @@ -43,6 +58,8 @@ class AzureAISearch extends Tool { fields.AZURE_AI_SEARCH_SERVICE_ENDPOINT, 'AZURE_AI_SEARCH_SERVICE_ENDPOINT', ); + this.isUserProvidedEndpoint = + fields.userProvidedAuthFields?.has('AZURE_AI_SEARCH_SERVICE_ENDPOINT') === true; this.indexName = this._initializeField( fields.AZURE_AI_SEARCH_INDEX_NAME, 'AZURE_AI_SEARCH_INDEX_NAME', @@ -75,6 +92,12 @@ class AzureAISearch extends Tool { ); } + if (this.isUserProvidedEndpoint && !isAzureSearchEndpoint(this.serviceEndpoint)) { + throw new Error( + 'User-provided Azure AI Search endpoints must use a trusted Azure Search host.', + ); + } + if (this.override) { return; } diff --git a/api/app/clients/tools/structured/AzureAISearch.spec.js b/api/app/clients/tools/structured/AzureAISearch.spec.js new file mode 100644 index 0000000000..abedd24a2a --- /dev/null +++ b/api/app/clients/tools/structured/AzureAISearch.spec.js @@ -0,0 +1,54 @@ +jest.mock( + '@azure/search-documents', + () => ({ + AzureKeyCredential: jest.fn(), + SearchClient: jest.fn(), + }), + { virtual: true }, +); + +jest.mock( + '@librechat/agents/langchain/tools', + () => ({ + Tool: class {}, + }), + { virtual: true }, +); + +jest.mock( + '@librechat/data-schemas', + () => ({ + logger: { error: jest.fn() }, + }), + { virtual: true }, +); + +const AzureAISearch = require('./AzureAISearch'); + +describe('AzureAISearch', () => { + const requiredFields = { + AZURE_AI_SEARCH_API_KEY: 'key', + AZURE_AI_SEARCH_INDEX_NAME: 'index', + userProvidedAuthFields: new Set(['AZURE_AI_SEARCH_SERVICE_ENDPOINT']), + }; + + it('rejects a user-provided endpoint outside Azure AI Search', () => { + expect( + () => + new AzureAISearch({ + ...requiredFields, + AZURE_AI_SEARCH_SERVICE_ENDPOINT: 'http://127.0.0.1:9000', + }), + ).toThrow('User-provided Azure AI Search endpoints must use a trusted Azure Search host.'); + }); + + it('accepts a user-provided Azure AI Search endpoint', () => { + expect( + () => + new AzureAISearch({ + ...requiredFields, + AZURE_AI_SEARCH_SERVICE_ENDPOINT: 'https://example.search.windows.net', + }), + ).not.toThrow(); + }); +}); diff --git a/api/app/clients/tools/structured/StableDiffusion.js b/api/app/clients/tools/structured/StableDiffusion.js index 89792a84b0..ac58bf6987 100644 --- a/api/app/clients/tools/structured/StableDiffusion.js +++ b/api/app/clients/tools/structured/StableDiffusion.js @@ -7,7 +7,7 @@ const { v4: uuidv4 } = require('uuid'); const { logger } = require('@librechat/data-schemas'); const { Tool } = require('@librechat/agents/langchain/tools'); const { FileContext, ContentTypes } = require('librechat-data-provider'); -const { getBasePath } = require('@librechat/api'); +const { applySSRFSafeAgentIfDirect, getBasePath } = require('@librechat/api'); const paths = require('~/config/paths'); const stableDiffusionJsonSchema = { @@ -54,6 +54,7 @@ class StableDiffusionAPI extends Tool { this.name = 'stable-diffusion'; this.url = fields.SD_WEBUI_URL || this.getServerURL(); + this.isUserProvidedEndpoint = fields.userProvidedAuthFields?.has('SD_WEBUI_URL') === true; this.description_for_model = `// Generate images and visuals using text. // Guidelines: // - ALWAYS use {{"prompt": "7+ detailed keywords", "negative_prompt": "7+ detailed keywords"}} structure for queries. @@ -116,7 +117,11 @@ class StableDiffusionAPI extends Tool { }; let generationResponse; try { - generationResponse = await axios.post(`${url}/sdapi/v1/txt2img`, payload); + const requestUrl = `${url}/sdapi/v1/txt2img`; + const requestConfig = this.isUserProvidedEndpoint + ? applySSRFSafeAgentIfDirect({}, requestUrl) + : undefined; + generationResponse = await axios.post(requestUrl, payload, requestConfig); } catch (error) { logger.error('[StableDiffusion] Error while generating image:', error); return this.returnValue('Error making API request.'); diff --git a/api/app/clients/tools/structured/StableDiffusion.spec.js b/api/app/clients/tools/structured/StableDiffusion.spec.js new file mode 100644 index 0000000000..f84d23c2f9 --- /dev/null +++ b/api/app/clients/tools/structured/StableDiffusion.spec.js @@ -0,0 +1,70 @@ +const axios = require('axios'); + +const mockApplySSRFSafeAgentIfDirect = jest.fn(); + +jest.mock('axios', () => ({ post: jest.fn() }), { virtual: true }); +jest.mock('fs'); +jest.mock('sharp', () => jest.fn(), { virtual: true }); +jest.mock('uuid', () => ({ v4: jest.fn() }), { virtual: true }); +jest.mock( + '@librechat/data-schemas', + () => ({ + logger: { error: jest.fn() }, + }), + { virtual: true }, +); +jest.mock( + '@librechat/agents/langchain/tools', + () => ({ + Tool: class {}, + }), + { virtual: true }, +); +jest.mock( + 'librechat-data-provider', + () => ({ + ContentTypes: {}, + FileContext: {}, + }), + { virtual: true }, +); +jest.mock( + '@librechat/api', + () => ({ + applySSRFSafeAgentIfDirect: (...args) => mockApplySSRFSafeAgentIfDirect(...args), + getBasePath: jest.fn(), + }), + { virtual: true }, +); +jest.mock('~/config/paths', () => ({}), { virtual: true }); + +const StableDiffusionAPI = require('./StableDiffusion'); + +describe('StableDiffusionAPI', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('uses a connect-time SSRF guard for a user-provided endpoint', async () => { + const error = new Error('SSRF protection: blocked address'); + mockApplySSRFSafeAgentIfDirect.mockImplementation(() => { + throw error; + }); + const tool = new StableDiffusionAPI({ + SD_WEBUI_URL: 'http://127.0.0.1:9000', + userProvidedAuthFields: new Set(['SD_WEBUI_URL']), + }); + + const result = await tool._call({ + prompt: 'test prompt', + negative_prompt: 'test negative', + }); + + expect(mockApplySSRFSafeAgentIfDirect).toHaveBeenCalledWith( + {}, + 'http://127.0.0.1:9000/sdapi/v1/txt2img', + ); + expect(axios.post).not.toHaveBeenCalled(); + expect(result).toBe('Error making API request.'); + }); +}); diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index e0f354e367..1eb28acd2c 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -23,6 +23,7 @@ const { resolveCodeExecutionContext, } = require('@librechat/api'); const { + AuthType, Tools, Constants, Permissions, @@ -142,7 +143,20 @@ const validateTools = async (user, tools = []) => { const loadToolWithAuth = (userId, authFields, ToolConstructor, options = {}) => { return async function () { const authValues = await loadAuthValues({ userId, authFields }); - return new ToolConstructor({ ...options, ...authValues, userId }); + const userProvidedAuthFields = new Set( + authFields + .flatMap((authField) => authField.split('||')) + .filter((authField) => { + const value = process.env[authField]; + return !value || value.trim() === '' || value === AuthType.USER_PROVIDED; + }), + ); + return new ToolConstructor({ + ...options, + ...authValues, + userId, + userProvidedAuthFields, + }); }; }; diff --git a/api/app/clients/tools/util/handleTools.test.js b/api/app/clients/tools/util/handleTools.test.js index 747ea9c2ba..63953a1ec3 100644 --- a/api/app/clients/tools/util/handleTools.test.js +++ b/api/app/clients/tools/util/handleTools.test.js @@ -310,6 +310,21 @@ describe('Tool Handlers', () => { expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(2); }); + it('marks credentials without an operator value as user-provided', async () => { + class CapturingTool { + constructor(fields) { + this.userProvidedAuthFields = fields.userProvidedAuthFields; + } + } + + process.env.SD_WEBUI_URL = 'user_provided'; + const initToolFunction = loadToolWithAuth('userId', ['SD_WEBUI_URL'], CapturingTool); + const tool = await initToolFunction(); + + expect(tool.userProvidedAuthFields).toEqual(new Set(['SD_WEBUI_URL'])); + delete process.env.SD_WEBUI_URL; + }); + it('should throw an error for an unauthenticated tool', async () => { try { await loadTool2();