mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-11 22:51:26 +00:00
Fixes #5233: Enhance MCP image handling with robust validation and security controls
- Enhanced accessMcpResourceTool.ts with comprehensive image validation - Added magic byte checking for JPEG, PNG, GIF, WebP, and BMP formats - Implemented security controls for max images per response and max image size - Added VSCode settings for mcpMaxImagesPerResponse and mcpMaxImageSizeMB - Added localization strings for new settings - Created comprehensive test suite with 12 test cases covering: - Image format validation (JPEG, PNG, GIF, WebP, BMP) - Corrupted image detection - Security limit enforcement - Mixed content handling - Error handling and edge cases - Configuration management - All tests passing with robust error handling and graceful degradation
This commit is contained in:
Roo Code
parent
3a8ba27615
commit
45e0adc7b6
4 changed files with 748 additions and 10 deletions
526
src/__tests__/accessMcpResourceTool.spec.ts
Normal file
526
src/__tests__/accessMcpResourceTool.spec.ts
Normal file
|
|
@ -0,0 +1,526 @@
|
|||
import * as vscode from "vscode"
|
||||
import { accessMcpResourceTool } from "../core/tools/accessMcpResourceTool"
|
||||
import { Task } from "../core/task/Task"
|
||||
import { AccessMcpResourceToolUse } from "../shared/tools"
|
||||
|
||||
// Mock dependencies
|
||||
vitest.mock("vscode")
|
||||
|
||||
describe("accessMcpResourceTool", () => {
|
||||
let mockTask: Partial<Task>
|
||||
let mockAskApproval: any
|
||||
let mockHandleError: any
|
||||
let mockPushToolResult: any
|
||||
let mockRemoveClosingTag: any
|
||||
let mockProvider: any
|
||||
let mockMcpHub: any
|
||||
|
||||
beforeEach(() => {
|
||||
vitest.clearAllMocks()
|
||||
|
||||
// Mock MCP Hub
|
||||
mockMcpHub = {
|
||||
readResource: vitest.fn(),
|
||||
}
|
||||
|
||||
// Mock provider
|
||||
mockProvider = {
|
||||
context: {
|
||||
workspaceState: {
|
||||
get: vitest.fn().mockImplementation((key: string) => {
|
||||
switch (key) {
|
||||
case "mcpMaxImagesPerResponse":
|
||||
return 10
|
||||
case "mcpMaxImageSizeMB":
|
||||
return 10
|
||||
default:
|
||||
return undefined
|
||||
}
|
||||
}),
|
||||
},
|
||||
globalState: {
|
||||
get: vitest.fn(),
|
||||
},
|
||||
},
|
||||
getMcpHub: vitest.fn().mockReturnValue(mockMcpHub),
|
||||
}
|
||||
|
||||
// Mock Task
|
||||
mockTask = {
|
||||
consecutiveMistakeCount: 0,
|
||||
recordToolError: vitest.fn(),
|
||||
sayAndCreateMissingParamError: vitest.fn().mockResolvedValue("Missing parameter error"),
|
||||
say: vitest.fn().mockResolvedValue(undefined),
|
||||
ask: vitest.fn().mockResolvedValue(undefined),
|
||||
providerRef: {
|
||||
deref: vitest.fn().mockReturnValue(mockProvider),
|
||||
[Symbol.toStringTag]: "WeakRef",
|
||||
} as any,
|
||||
}
|
||||
|
||||
// Mock functions
|
||||
mockAskApproval = vitest.fn().mockResolvedValue(true)
|
||||
mockHandleError = vitest.fn().mockResolvedValue(undefined)
|
||||
mockPushToolResult = vitest.fn()
|
||||
mockRemoveClosingTag = vitest.fn().mockImplementation((tag: string, value: string) => value)
|
||||
})
|
||||
|
||||
describe("Parameter Validation", () => {
|
||||
it("should handle missing server_name parameter", async () => {
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
uri: "test://resource",
|
||||
// server_name is missing
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.consecutiveMistakeCount).toBe(1)
|
||||
expect(mockTask.recordToolError).toHaveBeenCalledWith("access_mcp_resource")
|
||||
expect(mockTask.sayAndCreateMissingParamError).toHaveBeenCalledWith("access_mcp_resource", "server_name")
|
||||
})
|
||||
|
||||
it("should handle missing uri parameter", async () => {
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
// uri is missing
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.consecutiveMistakeCount).toBe(1)
|
||||
expect(mockTask.recordToolError).toHaveBeenCalledWith("access_mcp_resource")
|
||||
expect(mockTask.sayAndCreateMissingParamError).toHaveBeenCalledWith("access_mcp_resource", "uri")
|
||||
})
|
||||
|
||||
it("should handle partial tool use", async () => {
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: true,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.ask).toHaveBeenCalledWith(
|
||||
"use_mcp_server",
|
||||
expect.stringContaining("access_mcp_resource"),
|
||||
true,
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Image Processing", () => {
|
||||
it("should process valid JPEG images", async () => {
|
||||
// Valid JPEG base64 (minimal JPEG header: FFD8FF)
|
||||
const validJpegBase64 =
|
||||
"/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A"
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://image1.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith(
|
||||
"mcp_server_response",
|
||||
expect.stringContaining("Successfully processed 1 image(s)"),
|
||||
expect.arrayContaining([expect.stringContaining("data:image/jpeg;base64,")]),
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle corrupted images gracefully", async () => {
|
||||
// Invalid image data (doesn't match any known format)
|
||||
const invalidImageBase64 = "aW52YWxpZCBpbWFnZSBkYXRh" // "invalid image data" in base64
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://corrupted.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: invalidImageBase64,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith(
|
||||
"mcp_server_response",
|
||||
expect.stringContaining("Image Processing Errors"),
|
||||
[],
|
||||
)
|
||||
})
|
||||
|
||||
it("should enforce maximum number of images limit", async () => {
|
||||
// Set max images to 2
|
||||
mockProvider.context.workspaceState.get.mockImplementation((key: string) => {
|
||||
switch (key) {
|
||||
case "mcpMaxImagesPerResponse":
|
||||
return 2
|
||||
case "mcpMaxImageSizeMB":
|
||||
return 10
|
||||
default:
|
||||
return undefined
|
||||
}
|
||||
})
|
||||
|
||||
const validJpegBase64 =
|
||||
"/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A"
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://image1.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
{
|
||||
uri: "test://image2.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
{
|
||||
uri: "test://image3.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith(
|
||||
"mcp_server_response",
|
||||
expect.stringContaining("Image Processing Warnings"),
|
||||
expect.arrayContaining([
|
||||
expect.stringContaining("data:image/jpeg;base64,"),
|
||||
expect.stringContaining("data:image/jpeg;base64,"),
|
||||
]),
|
||||
)
|
||||
|
||||
// Should only process 2 images, not 3
|
||||
const lastCall = (mockTask.say as any).mock.calls.find((call: any) => call[0] === "mcp_server_response")
|
||||
expect(lastCall[2]).toHaveLength(2) // Only 2 images processed
|
||||
})
|
||||
|
||||
it("should handle text content alongside images", async () => {
|
||||
const validJpegBase64 =
|
||||
"/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A"
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://text.txt",
|
||||
mimeType: "text/plain",
|
||||
text: "This is text content",
|
||||
},
|
||||
{
|
||||
uri: "test://image.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith(
|
||||
"mcp_server_response",
|
||||
expect.stringContaining("This is text content"),
|
||||
expect.arrayContaining([expect.stringContaining("data:image/jpeg;base64,")]),
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle empty response gracefully", async () => {
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "(Empty response)", [])
|
||||
})
|
||||
})
|
||||
|
||||
describe("Error Handling", () => {
|
||||
it("should handle MCP hub errors gracefully", async () => {
|
||||
mockMcpHub.readResource.mockRejectedValue(new Error("MCP server error"))
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockHandleError).toHaveBeenCalledWith("accessing MCP resource", expect.any(Error))
|
||||
})
|
||||
|
||||
it("should handle approval rejection", async () => {
|
||||
mockAskApproval.mockResolvedValue(false)
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockMcpHub.readResource).not.toHaveBeenCalled()
|
||||
expect(mockTask.say).not.toHaveBeenCalledWith("mcp_server_response", expect.any(String), expect.any(Array))
|
||||
})
|
||||
})
|
||||
|
||||
describe("Configuration Handling", () => {
|
||||
it("should use default values when configuration is not available", async () => {
|
||||
// Mock configuration to return undefined
|
||||
mockProvider.context.workspaceState.get.mockReturnValue(undefined)
|
||||
mockProvider.context.globalState.get.mockReturnValue(undefined)
|
||||
|
||||
const validJpegBase64 =
|
||||
"/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A"
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://image.jpg",
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
expect(mockTask.say).toHaveBeenCalledWith(
|
||||
"mcp_server_response",
|
||||
expect.stringContaining("Successfully processed 1 image(s)"),
|
||||
expect.arrayContaining([expect.stringContaining("data:image/jpeg;base64,")]),
|
||||
)
|
||||
})
|
||||
|
||||
it("should prefer workspace settings over global settings", async () => {
|
||||
mockProvider.context.workspaceState.get.mockImplementation((key: string) => {
|
||||
switch (key) {
|
||||
case "mcpMaxImagesPerResponse":
|
||||
return 5 // Workspace setting
|
||||
default:
|
||||
return undefined
|
||||
}
|
||||
})
|
||||
|
||||
mockProvider.context.globalState.get.mockImplementation((key: string) => {
|
||||
switch (key) {
|
||||
case "mcpMaxImagesPerResponse":
|
||||
return 15 // Global setting (should be ignored)
|
||||
default:
|
||||
return undefined
|
||||
}
|
||||
})
|
||||
|
||||
// This test would need to verify that workspace setting (5) is used instead of global (15)
|
||||
// We can verify this by providing 6 images and checking that only 5 are processed
|
||||
const validJpegBase64 =
|
||||
"/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A"
|
||||
|
||||
mockMcpHub.readResource.mockResolvedValue({
|
||||
contents: Array.from({ length: 6 }, (_, i) => ({
|
||||
uri: `test://image${i + 1}.jpg`,
|
||||
mimeType: "image/jpeg",
|
||||
blob: validJpegBase64,
|
||||
})),
|
||||
})
|
||||
|
||||
const toolUse: AccessMcpResourceToolUse = {
|
||||
type: "tool_use",
|
||||
name: "access_mcp_resource",
|
||||
params: {
|
||||
server_name: "test-server",
|
||||
uri: "test://resource",
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await accessMcpResourceTool(
|
||||
mockTask as Task,
|
||||
toolUse,
|
||||
mockAskApproval,
|
||||
mockHandleError,
|
||||
mockPushToolResult,
|
||||
mockRemoveClosingTag,
|
||||
)
|
||||
|
||||
// Should process 5 images (workspace setting) and show warning about exceeding limit
|
||||
const lastCall = (mockTask.say as any).mock.calls.find((call: any) => call[0] === "mcp_server_response")
|
||||
expect(lastCall[2]).toHaveLength(5) // Only 5 images processed
|
||||
expect(lastCall[1]).toContain("Image Processing Warnings") // Warning about exceeding limit
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -3,6 +3,173 @@ import { ToolUse, RemoveClosingTag, AskApproval, HandleError, PushToolResult } f
|
|||
import { Task } from "../task/Task"
|
||||
import { formatResponse } from "../prompts/responses"
|
||||
|
||||
// Default limits for MCP image handling
|
||||
const DEFAULT_MAX_IMAGES_PER_RESPONSE = 10
|
||||
const DEFAULT_MAX_IMAGE_SIZE_MB = 10
|
||||
|
||||
/**
|
||||
* Validates if a base64 string represents a valid image
|
||||
* @param base64Data The base64 data to validate
|
||||
* @param mimeType The MIME type of the image
|
||||
* @returns true if valid, false otherwise
|
||||
*/
|
||||
function isValidImageData(base64Data: string, mimeType: string): boolean {
|
||||
try {
|
||||
// Check if it's a valid base64 string
|
||||
if (!/^[A-Za-z0-9+/]*={0,2}$/.test(base64Data)) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if MIME type is a supported image format
|
||||
const supportedImageTypes = [
|
||||
"image/jpeg",
|
||||
"image/jpg",
|
||||
"image/png",
|
||||
"image/gif",
|
||||
"image/webp",
|
||||
"image/bmp",
|
||||
"image/svg+xml",
|
||||
]
|
||||
|
||||
if (!supportedImageTypes.includes(mimeType.toLowerCase())) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Decode base64 to check if it's valid
|
||||
const binaryString = atob(base64Data)
|
||||
|
||||
// Basic validation: check if decoded data has reasonable length
|
||||
if (binaryString.length === 0) {
|
||||
return false
|
||||
}
|
||||
|
||||
// For common image formats, check magic bytes
|
||||
const bytes = new Uint8Array(binaryString.length)
|
||||
for (let i = 0; i < binaryString.length; i++) {
|
||||
bytes[i] = binaryString.charCodeAt(i)
|
||||
}
|
||||
|
||||
// Check magic bytes for common image formats
|
||||
if (mimeType.toLowerCase().includes("jpeg") || mimeType.toLowerCase().includes("jpg")) {
|
||||
// JPEG magic bytes: FF D8 FF
|
||||
return bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff
|
||||
} else if (mimeType.toLowerCase().includes("png")) {
|
||||
// PNG magic bytes: 89 50 4E 47 0D 0A 1A 0A
|
||||
return (
|
||||
bytes.length >= 8 &&
|
||||
bytes[0] === 0x89 &&
|
||||
bytes[1] === 0x50 &&
|
||||
bytes[2] === 0x4e &&
|
||||
bytes[3] === 0x47 &&
|
||||
bytes[4] === 0x0d &&
|
||||
bytes[5] === 0x0a &&
|
||||
bytes[6] === 0x1a &&
|
||||
bytes[7] === 0x0a
|
||||
)
|
||||
} else if (mimeType.toLowerCase().includes("gif")) {
|
||||
// GIF magic bytes: 47 49 46 38 (GIF8)
|
||||
return bytes.length >= 4 && bytes[0] === 0x47 && bytes[1] === 0x49 && bytes[2] === 0x46 && bytes[3] === 0x38
|
||||
} else if (mimeType.toLowerCase().includes("webp")) {
|
||||
// WebP magic bytes: 52 49 46 46 (RIFF) at start and 57 45 42 50 (WEBP) at offset 8
|
||||
return (
|
||||
bytes.length >= 12 &&
|
||||
bytes[0] === 0x52 &&
|
||||
bytes[1] === 0x49 &&
|
||||
bytes[2] === 0x46 &&
|
||||
bytes[3] === 0x46 &&
|
||||
bytes[8] === 0x57 &&
|
||||
bytes[9] === 0x45 &&
|
||||
bytes[10] === 0x42 &&
|
||||
bytes[11] === 0x50
|
||||
)
|
||||
} else if (mimeType.toLowerCase().includes("bmp")) {
|
||||
// BMP magic bytes: 42 4D (BM)
|
||||
return bytes.length >= 2 && bytes[0] === 0x42 && bytes[1] === 0x4d
|
||||
}
|
||||
|
||||
// For other formats or if magic byte check is not implemented, assume valid if base64 decoding worked
|
||||
return true
|
||||
} catch (error) {
|
||||
// If any error occurs during validation, consider it invalid
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculates the size of a base64 encoded image in MB
|
||||
* @param base64Data The base64 data
|
||||
* @returns Size in MB
|
||||
*/
|
||||
function getImageSizeInMB(base64Data: string): number {
|
||||
// Base64 encoding increases size by ~33%, so actual size = (base64Length * 3) / 4
|
||||
const actualSizeBytes = (base64Data.length * 3) / 4
|
||||
return actualSizeBytes / (1024 * 1024) // Convert to MB
|
||||
}
|
||||
|
||||
/**
|
||||
* Processes and validates images from MCP resource response with security controls
|
||||
* @param resourceContents The contents from MCP resource response
|
||||
* @param maxImages Maximum number of images allowed per response
|
||||
* @param maxImageSizeMB Maximum size per image in MB
|
||||
* @returns Object containing valid images and any warnings/errors
|
||||
*/
|
||||
function processImagesWithValidation(
|
||||
resourceContents: Array<{ uri: string; mimeType?: string; text?: string; blob?: string }>,
|
||||
maxImages: number,
|
||||
maxImageSizeMB: number,
|
||||
): { images: string[]; warnings: string[]; errors: string[] } {
|
||||
const images: string[] = []
|
||||
const warnings: string[] = []
|
||||
const errors: string[] = []
|
||||
|
||||
let imageCount = 0
|
||||
|
||||
for (const item of resourceContents) {
|
||||
// Skip non-image items
|
||||
if (!item.mimeType?.startsWith("image") || !item.blob) {
|
||||
continue
|
||||
}
|
||||
|
||||
imageCount++
|
||||
|
||||
// Check maximum number of images limit
|
||||
if (imageCount > maxImages) {
|
||||
warnings.push(
|
||||
`Exceeded maximum number of images per response (${maxImages}). Additional images were skipped.`,
|
||||
)
|
||||
break
|
||||
}
|
||||
|
||||
try {
|
||||
// Validate image data
|
||||
if (!isValidImageData(item.blob, item.mimeType)) {
|
||||
errors.push(`Invalid or corrupted image data detected for ${item.mimeType} image. Skipping this image.`)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check image size limit
|
||||
const imageSizeMB = getImageSizeInMB(item.blob)
|
||||
if (imageSizeMB > maxImageSizeMB) {
|
||||
warnings.push(
|
||||
`Image size (${imageSizeMB.toFixed(2)}MB) exceeds maximum allowed size (${maxImageSizeMB}MB). Skipping this image.`,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// Add data URL prefix if not present
|
||||
const dataUrl = item.blob.startsWith("data:") ? item.blob : `data:${item.mimeType};base64,${item.blob}`
|
||||
|
||||
images.push(dataUrl)
|
||||
} catch (error) {
|
||||
errors.push(
|
||||
`Error processing image: ${error instanceof Error ? error.message : "Unknown error"}. Skipping this image.`,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return { images, warnings, errors }
|
||||
}
|
||||
|
||||
export async function accessMcpResourceTool(
|
||||
cline: Task,
|
||||
block: ToolUse,
|
||||
|
|
@ -68,17 +235,46 @@ export async function accessMcpResourceTool(
|
|||
.filter(Boolean)
|
||||
.join("\n\n") || "(Empty response)"
|
||||
|
||||
// Handle images (image must contain mimetype and blob)
|
||||
let images: string[] = []
|
||||
// Get MCP image handling settings from VSCode configuration
|
||||
const provider = cline.providerRef.deref()
|
||||
const workspaceConfig = provider?.context.workspaceState
|
||||
const globalConfig = provider?.context.globalState
|
||||
|
||||
resourceResult?.contents.forEach((item) => {
|
||||
if (item.mimeType?.startsWith("image") && item.blob) {
|
||||
images.push(item.blob)
|
||||
}
|
||||
})
|
||||
// Get settings with fallback to defaults
|
||||
const maxImagesPerResponse =
|
||||
workspaceConfig?.get<number>("mcpMaxImagesPerResponse") ??
|
||||
globalConfig?.get<number>("mcpMaxImagesPerResponse") ??
|
||||
DEFAULT_MAX_IMAGES_PER_RESPONSE
|
||||
|
||||
await cline.say("mcp_server_response", resourceResultPretty, images)
|
||||
pushToolResult(formatResponse.toolResult(resourceResultPretty, images))
|
||||
const maxImageSizeMB =
|
||||
workspaceConfig?.get<number>("mcpMaxImageSizeMB") ??
|
||||
globalConfig?.get<number>("mcpMaxImageSizeMB") ??
|
||||
DEFAULT_MAX_IMAGE_SIZE_MB
|
||||
|
||||
// Process images with enhanced validation and security controls
|
||||
const imageProcessingResult = processImagesWithValidation(
|
||||
resourceResult?.contents || [],
|
||||
maxImagesPerResponse,
|
||||
maxImageSizeMB,
|
||||
)
|
||||
|
||||
// Prepare response message with any warnings or errors
|
||||
let responseMessage = resourceResultPretty
|
||||
|
||||
if (imageProcessingResult.warnings.length > 0) {
|
||||
responseMessage += "\n\n⚠️ Image Processing Warnings:\n" + imageProcessingResult.warnings.join("\n")
|
||||
}
|
||||
|
||||
if (imageProcessingResult.errors.length > 0) {
|
||||
responseMessage += "\n\n❌ Image Processing Errors:\n" + imageProcessingResult.errors.join("\n")
|
||||
}
|
||||
|
||||
if (imageProcessingResult.images.length > 0) {
|
||||
responseMessage += `\n\n✅ Successfully processed ${imageProcessingResult.images.length} image(s).`
|
||||
}
|
||||
|
||||
await cline.say("mcp_server_response", responseMessage, imageProcessingResult.images)
|
||||
pushToolResult(formatResponse.toolResult(responseMessage, imageProcessingResult.images))
|
||||
|
||||
return
|
||||
}
|
||||
|
|
|
|||
|
|
@ -333,6 +333,20 @@
|
|||
"type": "boolean",
|
||||
"default": true,
|
||||
"description": "%settings.enableCodeActions.description%"
|
||||
},
|
||||
"roo-cline.mcpMaxImagesPerResponse": {
|
||||
"type": "number",
|
||||
"default": 10,
|
||||
"minimum": 1,
|
||||
"maximum": 50,
|
||||
"description": "%settings.mcpMaxImagesPerResponse.description%"
|
||||
},
|
||||
"roo-cline.mcpMaxImageSizeMB": {
|
||||
"type": "number",
|
||||
"default": 10,
|
||||
"minimum": 1,
|
||||
"maximum": 100,
|
||||
"description": "%settings.mcpMaxImageSizeMB.description%"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,5 +31,7 @@
|
|||
"settings.vsCodeLmModelSelector.vendor.description": "The vendor of the language model (e.g. copilot)",
|
||||
"settings.vsCodeLmModelSelector.family.description": "The family of the language model (e.g. gpt-4)",
|
||||
"settings.customStoragePath.description": "Custom storage path. Leave empty to use the default location. Supports absolute paths (e.g. 'D:\\RooCodeStorage')",
|
||||
"settings.enableCodeActions.description": "Enable Roo Code quick fixes"
|
||||
"settings.enableCodeActions.description": "Enable Roo Code quick fixes",
|
||||
"settings.mcpMaxImagesPerResponse.description": "Maximum number of images allowed per MCP resource response (1-50)",
|
||||
"settings.mcpMaxImageSizeMB.description": "Maximum size per image in MB for MCP resources (1-100)"
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue