fix(databricks): accept DoD, custom-URL, and trailing-dot workspace hosts

Author: waleedlatif1

apps/sim/lib/core/security/input-validation.ts

@@ -1167,7 +1167,8 @@ function validateVendorHostedUrl(
   if (!urlResult.isValid) return urlResult
 
   const parsed = new URL(candidate)
-  const hostname = parsed.hostname.toLowerCase()
+  // A fully qualified name's trailing dot names the same host.
+  const hostname = parsed.hostname.toLowerCase().replace(/\.$/, '')
   const allowed = suffixes.some(
     (suffix) => (allowBareSuffix && hostname === suffix.slice(1)) || hostname.endsWith(suffix)
   )
@@ -1267,15 +1268,20 @@ export function validateWorkdayTenantUrl(
 }
 
 /**
- * Every production Databricks control-plane DNS zone, mirroring `ALL_ENVS` in the
- * Databricks SDK (`databricks/sdk/environments.py`). The SDK's `.dev.*`/`.staging.*`
- * zones are internal and deliberately omitted; the ones that are subdomains of a
- * zone listed here (e.g. `.staging.cloud.databricks.com`) match by suffix anyway.
+ * Every production Databricks control-plane DNS zone. All but `.cloud.databricks.mil` mirror
+ * `ALL_ENVS` in the Databricks SDK (`databricks/sdk/environments.py`); the DoD zone comes from the
+ * Databricks AWS GovCloud docs. The SDK's `.dev.*`/`.staging.*` zones are internal and
+ * deliberately omitted; the ones that are subdomains of a zone listed here (e.g.
+ * `.staging.cloud.databricks.com`) match by suffix anyway. `.databricks.com` admits workspace
+ * custom URLs (`acme.databricks.com`) and subsumes the AWS and GCP zones, which stay listed so
+ * the rejection message names them.
  */
 const DATABRICKS_ALLOWED_HOST_SUFFIXES = [
   '.cloud.databricks.com',
   '.cloud.databricks.us',
+  '.cloud.databricks.mil',
   '.gcp.databricks.com',
+  '.databricks.com',
   '.azuredatabricks.net',
   '.databricks.azure.us',
   '.databricks.azure.cn',
@@ -1288,6 +1294,8 @@ const DATABRICKS_ALLOWED_HOST_SUFFIXES = [
  * every REST call is made against it. Example valid hosts:
  * - dbc-1234abcd-5678.cloud.databricks.com (AWS)
  * - dbc-1234abcd-5678.cloud.databricks.us (AWS GovCloud)
+ * - dbc-1234abcd-5678.cloud.databricks.mil (AWS GovCloud DoD)
+ * - acme.databricks.com (workspace custom URL)
  * - adb-1234567890123456.7.azuredatabricks.net (Azure)
  * - adb-1234567890123456.7.databricks.azure.us (Azure US Government)
  * - adb-1234567890123456.7.databricks.azure.cn (Azure China)

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

@@ -1,17 +1,16 @@
-import { inputValidationMock, inputValidationMockFns } from '@sim/testing'
+import { inputValidationMock } from '@sim/testing'
+import { partialToolRegistry } from '@sim/testing/mocks/tool-registry.mock'
 import { describe, expect, it, vi } from 'vitest'
 
 vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
 
-/** Only this service's configs are needed; the full registry is ~6,000 modules. */
-vi.mock('@/tools/registry', async () => {
-  const { partialToolRegistry } = await import('@sim/testing/mocks/tool-registry.mock')
-  return { tools: partialToolRegistry(await import('@/tools/databricks')) }
-})
-
+import * as databricksTools from '@/tools/databricks'
 import { executeTool } from '@/tools/index'
 import { tools } from '@/tools/registry'
 
+/** Registers only this service's configs in the global registry mock; the full one is ~6,000 modules. */
+Object.assign(tools, partialToolRegistry(databricksTools))
+
 type ToolParams = Record<string, unknown>
 
 function urlBuilder(toolId: string): (params: ToolParams) => string {
@@ -40,15 +39,19 @@ const REQUEST_PARAMS = {
 const DATABRICKS_TOOL_IDS = Object.keys(tools).filter((id) => id.startsWith('databricks_'))
 
 describe('databricks workspace host allowlist', () => {
+  it('registers the Databricks tools', () => {
+    expect(DATABRICKS_TOOL_IDS).not.toEqual([])
+  })
+
   it.each([
     'attacker.example.com',
     'https://attacker.example.com/',
     'dbc-1.cloud.databricks.com.attacker.example.com',
     'attacker.example.com/dbc-1.cloud.databricks.com',
     'dbc-1.cloud.databricks.com@attacker.example.com',
-    'cloud.databricks.com',
+    'databricks.com',
+    'acme-databricks.com',
   ])('refuses %s in every tool before a URL is built', (host) => {
-    expect(DATABRICKS_TOOL_IDS.length).toBe(26)
     const accepted = DATABRICKS_TOOL_IDS.filter((id) => {
       try {
         urlBuilder(id)({ ...REQUEST_PARAMS, host })
@@ -60,16 +63,14 @@ describe('databricks workspace host allowlist', () => {
     expect(accepted).toEqual([])
   })
 
-  it('fails the tool call without sending the token to a foreign host', async () => {
+  it('fails the tool call for a foreign host with the allowlist error', async () => {
     const result = await executeTool('databricks_list_clusters', {
       host: 'attacker.example.com',
       apiKey: 'dapi-test-token',
     })
 
     expect(result.success).toBe(false)
     expect(result.error).toContain('host must be a Databricks-hosted domain')
-    expect(inputValidationMockFns.mockValidateUrlWithDNS).not.toHaveBeenCalled()
-    expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).not.toHaveBeenCalled()
   })
 
   it.each([
@@ -81,6 +82,9 @@ describe('databricks workspace host allowlist', () => {
     ['dbc-a1b2.cloud.databricks.us', 'https://dbc-a1b2.cloud.databricks.us'],
     ['adb-123.4.databricks.azure.us', 'https://adb-123.4.databricks.azure.us'],
     ['adb-123.4.databricks.azure.cn', 'https://adb-123.4.databricks.azure.cn'],
+    ['dbc-a1b2.cloud.databricks.mil', 'https://dbc-a1b2.cloud.databricks.mil'],
+    ['https://acme.databricks.com/', 'https://acme.databricks.com'],
+    ['dbc-a1b2.cloud.databricks.com.', 'https://dbc-a1b2.cloud.databricks.com.'],
   ])('builds the same request URLs for workspace host %s', (host, origin) => {
     const params = { ...REQUEST_PARAMS, host }
     expect(urlBuilder('databricks_list_clusters')(params)).toBe(`${origin}/api/2.0/clusters/list`)