|
| 1 | +import { inputValidationMock } from '@sim/testing' |
| 2 | +import { partialToolRegistry } from '@sim/testing/mocks/tool-registry.mock' |
| 3 | +import { getErrorMessage } from '@sim/utils/errors' |
| 4 | +import { describe, expect, it, vi } from 'vitest' |
| 5 | + |
| 6 | +vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock) |
| 7 | + |
| 8 | +import * as databricksTools from '@/tools/databricks' |
| 9 | +import { executeTool } from '@/tools/index' |
| 10 | +import { tools } from '@/tools/registry' |
| 11 | + |
| 12 | +/** Registers only this service's configs in the global registry mock; the full one is ~6,000 modules. */ |
| 13 | +Object.assign(tools, partialToolRegistry(databricksTools)) |
| 14 | + |
| 15 | +type ToolParams = Record<string, unknown> |
| 16 | + |
| 17 | +function urlBuilder(toolId: string): (params: ToolParams) => string { |
| 18 | + const url = tools[toolId].request.url |
| 19 | + if (typeof url !== 'function') throw new Error(`${toolId} has a static url`) |
| 20 | + return url as (params: ToolParams) => string |
| 21 | +} |
| 22 | + |
| 23 | +/** Every identifier any Databricks tool reads while building its URL. */ |
| 24 | +const REQUEST_PARAMS = { |
| 25 | + apiKey: 'dapi-test-token', |
| 26 | + spaceId: 'space1', |
| 27 | + conversationId: 'conv1', |
| 28 | + messageId: 'msg1', |
| 29 | + attachmentId: 'att1', |
| 30 | + statementId: 'stmt1', |
| 31 | + clusterId: 'cluster1', |
| 32 | + jobId: 1, |
| 33 | + runId: 2, |
| 34 | + warehouseId: 'wh1', |
| 35 | + content: 'question', |
| 36 | + sql: 'SELECT 1', |
| 37 | + rating: 'POSITIVE', |
| 38 | +} |
| 39 | + |
| 40 | +/** The validator's refusal, naming the `host` param and every allowlisted Databricks domain. */ |
| 41 | +const HOST_ALLOWLIST_ERROR = |
| 42 | + 'host must be a Databricks-hosted domain (e.g., *.cloud.databricks.com, *.cloud.databricks.us, *.cloud.databricks.mil, *.gcp.databricks.com, *.databricks.com, *.azuredatabricks.net, *.databricks.azure.us, *.databricks.azure.cn)' |
| 43 | + |
| 44 | +const DATABRICKS_TOOL_IDS = Object.keys(tools).filter((id) => id.startsWith('databricks_')) |
| 45 | + |
| 46 | +describe('databricks workspace host allowlist', () => { |
| 47 | + it.each([ |
| 48 | + 'attacker.example.com', |
| 49 | + 'https://attacker.example.com/', |
| 50 | + 'dbc-1.cloud.databricks.com.attacker.example.com', |
| 51 | + 'attacker.example.com/dbc-1.cloud.databricks.com', |
| 52 | + 'dbc-1.cloud.databricks.com@attacker.example.com', |
| 53 | + 'databricks.com', |
| 54 | + 'acme-databricks.com', |
| 55 | + ])('refuses %s in every tool with the host allowlist error', (host) => { |
| 56 | + const notRefusedByAllowlist = DATABRICKS_TOOL_IDS.map((id) => { |
| 57 | + try { |
| 58 | + return `${id}: built ${urlBuilder(id)({ ...REQUEST_PARAMS, host })}` |
| 59 | + } catch (error) { |
| 60 | + return `${id}: ${getErrorMessage(error)}` |
| 61 | + } |
| 62 | + }).filter((outcome) => !outcome.endsWith(`: ${HOST_ALLOWLIST_ERROR}`)) |
| 63 | + expect(notRefusedByAllowlist).toEqual([]) |
| 64 | + }) |
| 65 | + |
| 66 | + it('fails the tool call for a foreign host with the allowlist error', async () => { |
| 67 | + const result = await executeTool('databricks_list_clusters', { |
| 68 | + host: 'attacker.example.com', |
| 69 | + apiKey: 'dapi-test-token', |
| 70 | + }) |
| 71 | + |
| 72 | + expect(result.success).toBe(false) |
| 73 | + expect(result.error).toContain(HOST_ALLOWLIST_ERROR) |
| 74 | + }) |
| 75 | + |
| 76 | + it.each([ |
| 77 | + ['dbc-a1b2.cloud.databricks.com', 'https://dbc-a1b2.cloud.databricks.com'], |
| 78 | + [' https://dbc-a1b2.cloud.databricks.com/ ', 'https://dbc-a1b2.cloud.databricks.com'], |
| 79 | + ['http://dbc-a1b2.cloud.databricks.com', 'https://dbc-a1b2.cloud.databricks.com'], |
| 80 | + ['adb-123.4.azuredatabricks.net', 'https://adb-123.4.azuredatabricks.net'], |
| 81 | + ['https://123.4.gcp.databricks.com/', 'https://123.4.gcp.databricks.com'], |
| 82 | + ['dbc-a1b2.cloud.databricks.us', 'https://dbc-a1b2.cloud.databricks.us'], |
| 83 | + ['adb-123.4.databricks.azure.us', 'https://adb-123.4.databricks.azure.us'], |
| 84 | + ['adb-123.4.databricks.azure.cn', 'https://adb-123.4.databricks.azure.cn'], |
| 85 | + ['dbc-a1b2.cloud.databricks.mil', 'https://dbc-a1b2.cloud.databricks.mil'], |
| 86 | + ['https://acme.databricks.com/', 'https://acme.databricks.com'], |
| 87 | + ])('builds the same request URLs for workspace host %s', (host, origin) => { |
| 88 | + const params = { ...REQUEST_PARAMS, host } |
| 89 | + expect(urlBuilder('databricks_list_clusters')(params)).toBe(`${origin}/api/2.0/clusters/list`) |
| 90 | + expect(urlBuilder('databricks_execute_sql')(params)).toBe(`${origin}/api/2.0/sql/statements/`) |
| 91 | + expect(urlBuilder('databricks_get_job')(params)).toBe(`${origin}/api/2.1/jobs/get?job_id=1`) |
| 92 | + expect(urlBuilder('databricks_get_run_output')(params)).toBe( |
| 93 | + `${origin}/api/2.1/jobs/runs/get-output?run_id=2` |
| 94 | + ) |
| 95 | + expect(urlBuilder('databricks_genie_get_message')(params)).toBe( |
| 96 | + `${origin}/api/2.0/genie/spaces/space1/conversations/conv1/messages/msg1` |
| 97 | + ) |
| 98 | + }) |
| 99 | + |
| 100 | + it.each([ |
| 101 | + ['dbc-a1b2.cloud.databricks.com.', 'https://dbc-a1b2.cloud.databricks.com'], |
| 102 | + ['https://dbc-a1b2.cloud.databricks.com.:8443/', 'https://dbc-a1b2.cloud.databricks.com:8443'], |
| 103 | + ])( |
| 104 | + 'drops the trailing FQDN dot of %s, which the old tools kept and Bun TLS rejects', |
| 105 | + (host, origin) => { |
| 106 | + expect(urlBuilder('databricks_list_clusters')({ ...REQUEST_PARAMS, host })).toBe( |
| 107 | + `${origin}/api/2.0/clusters/list` |
| 108 | + ) |
| 109 | + } |
| 110 | + ) |
| 111 | +}) |
0 commit comments