Skip to content

Commit d447fc7

Browse files
waleedlatif1waleedlatif1
committed
fix(databricks): send workspace tokens only to Databricks workspace hosts
1 parent 5c446f1 commit d447fc7

14 files changed

Lines changed: 133 additions & 80 deletions

‎apps/sim/tools/databricks/cancel_run.ts‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type {
22
DatabricksCancelRunParams,
33
DatabricksCancelRunResponse,
44
} from '@/tools/databricks/types'
5+
import { databricksUrl } from '@/tools/databricks/utils'
56
import type { ToolConfig } from '@/tools/types'
67

78
export const cancelRunTool: ToolConfig<DatabricksCancelRunParams, DatabricksCancelRunResponse> = {
@@ -33,13 +34,7 @@ export const cancelRunTool: ToolConfig<DatabricksCancelRunParams, DatabricksCanc
3334
},
3435

3536
request: {
36-
url: (params) => {
37-
const host = params.host
38-
.trim()
39-
.replace(/^https?:\/\//, '')
40-
.replace(/\/$/, '')
41-
return `https://${host}/api/2.1/jobs/runs/cancel`
42-
},
37+
url: (params) => databricksUrl(params.host, '/api/2.1/jobs/runs/cancel'),
4338
method: 'POST',
4439
headers: (params) => ({
4540
'Content-Type': 'application/json',

‎apps/sim/tools/databricks/databricks.test.ts‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import { inputValidationMock, inputValidationMockFns } from '@sim/testing'
2+
import { describe, expect, it, vi } from 'vitest'
3+
4+
vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
5+
6+
/** Only this service's configs are needed; the full registry is ~6,000 modules. */
7+
vi.mock('@/tools/registry', async () => {
8+
const { partialToolRegistry } = await import('@sim/testing/mocks/tool-registry.mock')
9+
return { tools: partialToolRegistry(await import('@/tools/databricks')) }
10+
})
11+
12+
import { executeTool } from '@/tools/index'
13+
import { tools } from '@/tools/registry'
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+
const DATABRICKS_TOOL_IDS = Object.keys(tools).filter((id) => id.startsWith('databricks_'))
41+
42+
describe('databricks workspace host allowlist', () => {
43+
it.each([
44+
'attacker.example.com',
45+
'https://attacker.example.com/',
46+
'dbc-1.cloud.databricks.com.attacker.example.com',
47+
'attacker.example.com/dbc-1.cloud.databricks.com',
48+
'dbc-1.cloud.databricks.com@attacker.example.com',
49+
'cloud.databricks.com',
50+
])('refuses %s in every tool before a URL is built', (host) => {
51+
expect(DATABRICKS_TOOL_IDS.length).toBe(26)
52+
const accepted = DATABRICKS_TOOL_IDS.filter((id) => {
53+
try {
54+
urlBuilder(id)({ ...REQUEST_PARAMS, host })
55+
return true
56+
} catch {
57+
return false
58+
}
59+
})
60+
expect(accepted).toEqual([])
61+
})
62+
63+
it('fails the tool call without sending the token to a foreign host', async () => {
64+
const result = await executeTool('databricks_list_clusters', {
65+
host: 'attacker.example.com',
66+
apiKey: 'dapi-test-token',
67+
})
68+
69+
expect(result.success).toBe(false)
70+
expect(result.error).toContain('host must be a Databricks-hosted domain')
71+
expect(inputValidationMockFns.mockValidateUrlWithDNS).not.toHaveBeenCalled()
72+
expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).not.toHaveBeenCalled()
73+
})
74+
75+
it.each([
76+
['dbc-a1b2.cloud.databricks.com', 'https://dbc-a1b2.cloud.databricks.com'],
77+
[' https://dbc-a1b2.cloud.databricks.com/ ', 'https://dbc-a1b2.cloud.databricks.com'],
78+
['http://dbc-a1b2.cloud.databricks.com', 'https://dbc-a1b2.cloud.databricks.com'],
79+
['adb-123.4.azuredatabricks.net', 'https://adb-123.4.azuredatabricks.net'],
80+
['https://123.4.gcp.databricks.com/', 'https://123.4.gcp.databricks.com'],
81+
['dbc-a1b2.cloud.databricks.us', 'https://dbc-a1b2.cloud.databricks.us'],
82+
['adb-123.4.databricks.azure.us', 'https://adb-123.4.databricks.azure.us'],
83+
['adb-123.4.databricks.azure.cn', 'https://adb-123.4.databricks.azure.cn'],
84+
])('builds the same request URLs for workspace host %s', (host, origin) => {
85+
const params = { ...REQUEST_PARAMS, host }
86+
expect(urlBuilder('databricks_list_clusters')(params)).toBe(`${origin}/api/2.0/clusters/list`)
87+
expect(urlBuilder('databricks_execute_sql')(params)).toBe(`${origin}/api/2.0/sql/statements/`)
88+
expect(urlBuilder('databricks_get_job')(params)).toBe(`${origin}/api/2.1/jobs/get?job_id=1`)
89+
expect(urlBuilder('databricks_get_run_output')(params)).toBe(
90+
`${origin}/api/2.1/jobs/runs/get-output?run_id=2`
91+
)
92+
expect(urlBuilder('databricks_genie_get_message')(params)).toBe(
93+
`${origin}/api/2.0/genie/spaces/space1/conversations/conv1/messages/msg1`
94+
)
95+
})
96+
})

‎apps/sim/tools/databricks/execute_sql.ts‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type {
22
DatabricksExecuteSqlParams,
33
DatabricksExecuteSqlResponse,
44
} from '@/tools/databricks/types'
5+
import { databricksUrl } from '@/tools/databricks/utils'
56
import type { ToolConfig } from '@/tools/types'
67

78
export const executeSqlTool: ToolConfig<DatabricksExecuteSqlParams, DatabricksExecuteSqlResponse> =
@@ -65,13 +66,7 @@ export const executeSqlTool: ToolConfig<DatabricksExecuteSqlParams, DatabricksEx
6566
},
6667

6768
request: {
68-
url: (params) => {
69-
const host = params.host
70-
.trim()
71-
.replace(/^https?:\/\//, '')
72-
.replace(/\/$/, '')
73-
return `https://${host}/api/2.0/sql/statements/`
74-
},
69+
url: (params) => databricksUrl(params.host, '/api/2.0/sql/statements/'),
7570
method: 'POST',
7671
headers: (params) => ({
7772
'Content-Type': 'application/json',

‎apps/sim/tools/databricks/get_cluster.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type {
22
DatabricksGetClusterParams,
33
DatabricksGetClusterResponse,
44
} from '@/tools/databricks/types'
5+
import { databricksUrl } from '@/tools/databricks/utils'
56
import type { ToolConfig } from '@/tools/types'
67

78
export const getClusterTool: ToolConfig<DatabricksGetClusterParams, DatabricksGetClusterResponse> =
@@ -35,11 +36,7 @@ export const getClusterTool: ToolConfig<DatabricksGetClusterParams, DatabricksGe
3536

3637
request: {
3738
url: (params) => {
38-
const host = params.host
39-
.trim()
40-
.replace(/^https?:\/\//, '')
41-
.replace(/\/$/, '')
42-
const url = new URL(`https://${host}/api/2.0/clusters/get`)
39+
const url = new URL(databricksUrl(params.host, '/api/2.0/clusters/get'))
4340
url.searchParams.set('cluster_id', params.clusterId.trim())
4441
return url.toString()
4542
},

‎apps/sim/tools/databricks/get_job.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { DatabricksGetJobParams, DatabricksGetJobResponse } from '@/tools/databricks/types'
2+
import { databricksUrl } from '@/tools/databricks/utils'
23
import type { ToolConfig } from '@/tools/types'
34

45
export const getJobTool: ToolConfig<DatabricksGetJobParams, DatabricksGetJobResponse> = {
@@ -30,11 +31,7 @@ export const getJobTool: ToolConfig<DatabricksGetJobParams, DatabricksGetJobResp
3031

3132
request: {
3233
url: (params) => {
33-
const host = params.host
34-
.trim()
35-
.replace(/^https?:\/\//, '')
36-
.replace(/\/$/, '')
37-
const url = new URL(`https://${host}/api/2.1/jobs/get`)
34+
const url = new URL(databricksUrl(params.host, '/api/2.1/jobs/get'))
3835
url.searchParams.set('job_id', String(params.jobId))
3936
return url.toString()
4037
},

‎apps/sim/tools/databricks/get_run.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { DatabricksGetRunParams, DatabricksGetRunResponse } from '@/tools/databricks/types'
2+
import { databricksUrl } from '@/tools/databricks/utils'
23
import type { ToolConfig } from '@/tools/types'
34

45
export const getRunTool: ToolConfig<DatabricksGetRunParams, DatabricksGetRunResponse> = {
@@ -42,11 +43,7 @@ export const getRunTool: ToolConfig<DatabricksGetRunParams, DatabricksGetRunResp
4243

4344
request: {
4445
url: (params) => {
45-
const host = params.host
46-
.trim()
47-
.replace(/^https?:\/\//, '')
48-
.replace(/\/$/, '')
49-
const url = new URL(`https://${host}/api/2.1/jobs/runs/get`)
46+
const url = new URL(databricksUrl(params.host, '/api/2.1/jobs/runs/get'))
5047
url.searchParams.set('run_id', String(params.runId))
5148
if (params.includeHistory) url.searchParams.set('include_history', 'true')
5249
if (params.includeResolvedValues) url.searchParams.set('include_resolved_values', 'true')

‎apps/sim/tools/databricks/get_run_output.ts‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type {
22
DatabricksGetRunOutputParams,
33
DatabricksGetRunOutputResponse,
44
} from '@/tools/databricks/types'
5+
import { databricksUrl } from '@/tools/databricks/utils'
56
import type { ToolConfig } from '@/tools/types'
67

78
export const getRunOutputTool: ToolConfig<
@@ -36,13 +37,8 @@ export const getRunOutputTool: ToolConfig<
3637
},
3738

3839
request: {
39-
url: (params) => {
40-
const host = params.host
41-
.trim()
42-
.replace(/^https?:\/\//, '')
43-
.replace(/\/$/, '')
44-
return `https://${host}/api/2.1/jobs/runs/get-output?run_id=${params.runId}`
45-
},
40+
url: (params) =>
41+
databricksUrl(params.host, `/api/2.1/jobs/runs/get-output?run_id=${params.runId}`),
4642
method: 'GET',
4743
headers: (params) => ({
4844
Accept: 'application/json',

‎apps/sim/tools/databricks/get_statement.ts‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type {
22
DatabricksExecuteSqlResponse,
33
DatabricksGetStatementParams,
44
} from '@/tools/databricks/types'
5+
import { databricksUrl } from '@/tools/databricks/utils'
56
import type { ToolConfig } from '@/tools/types'
67

78
export const getStatementTool: ToolConfig<
@@ -36,13 +37,8 @@ export const getStatementTool: ToolConfig<
3637
},
3738

3839
request: {
39-
url: (params) => {
40-
const host = params.host
41-
.trim()
42-
.replace(/^https?:\/\//, '')
43-
.replace(/\/$/, '')
44-
return `https://${host}/api/2.0/sql/statements/${params.statementId.trim()}`
45-
},
40+
url: (params) =>
41+
databricksUrl(params.host, `/api/2.0/sql/statements/${params.statementId.trim()}`),
4642
method: 'GET',
4743
headers: (params) => ({
4844
Accept: 'application/json',

‎apps/sim/tools/databricks/list_clusters.ts‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { DatabricksBaseParams, DatabricksListClustersResponse } from '@/tools/databricks/types'
2+
import { databricksUrl } from '@/tools/databricks/utils'
23
import type { ToolConfig } from '@/tools/types'
34

45
export const listClustersTool: ToolConfig<DatabricksBaseParams, DatabricksListClustersResponse> = {
@@ -24,13 +25,7 @@ export const listClustersTool: ToolConfig<DatabricksBaseParams, DatabricksListCl
2425
},
2526

2627
request: {
27-
url: (params) => {
28-
const host = params.host
29-
.trim()
30-
.replace(/^https?:\/\//, '')
31-
.replace(/\/$/, '')
32-
return `https://${host}/api/2.0/clusters/list`
33-
},
28+
url: (params) => databricksUrl(params.host, '/api/2.0/clusters/list'),
3429
method: 'GET',
3530
headers: (params) => ({
3631
Accept: 'application/json',

‎apps/sim/tools/databricks/list_jobs.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { DatabricksListJobsParams, DatabricksListJobsResponse } from '@/tools/databricks/types'
2+
import { databricksUrl } from '@/tools/databricks/utils'
23
import type { ToolConfig } from '@/tools/types'
34

45
export const listJobsTool: ToolConfig<DatabricksListJobsParams, DatabricksListJobsResponse> = {
@@ -48,11 +49,7 @@ export const listJobsTool: ToolConfig<DatabricksListJobsParams, DatabricksListJo
4849

4950
request: {
5051
url: (params) => {
51-
const host = params.host
52-
.trim()
53-
.replace(/^https?:\/\//, '')
54-
.replace(/\/$/, '')
55-
const url = new URL(`https://${host}/api/2.1/jobs/list`)
52+
const url = new URL(databricksUrl(params.host, '/api/2.1/jobs/list'))
5653
if (params.limit) url.searchParams.set('limit', String(params.limit))
5754
if (params.offset) url.searchParams.set('offset', String(params.offset))
5855
if (params.name) url.searchParams.set('name', params.name)

0 commit comments

Comments
 (0)