fix(snowflake): verify key pairs through the egress guard and leave an organization reconnect's mode unchosen

Key-pair validation now sends its signed JWT through the same guarded egress
path as keyless, so an account host that resolves to a private address is
refused as host_not_allowed. An organization credential's stored mode is not
reported, so its reconnect starts with no mode chosen instead of defaulting to a
programmatic access token. The Infisical and ServiceNow grant tests use the
shared jsonResponse helper.

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/integrations/components/connect-service-account-modal/snowflake-account-modal.tsx

@@ -156,14 +156,15 @@ function SnowflakeAccountForm({
   const [displayName, setDisplayName] = useState(initialDisplayName ?? '')
   const [description, setDescription] = useState(initialDescription ?? '')
 
-  // A reconnect opens on the credential's stored mode, so it never switches auth by default;
-  // a new connection defaults to keyless wherever this deployment supports it.
-  const defaultAuthType: SnowflakeAuthType = credentialId
-    ? (keyless.reconnectAuthType ?? 'programmatic_access_token')
+  // A reconnect opens on the credential's stored mode, so it never switches auth by default. An
+  // organization credential's mode is not reported, so its reconnect starts with no mode chosen
+  // rather than a guess. A new connection defaults to keyless wherever this deployment supports it.
+  const defaultAuthType: SnowflakeAuthType | null = credentialId
+    ? (keyless.reconnectAuthType ?? null)
     : keyless.available
       ? 'workload_identity'
       : 'programmatic_access_token'
-  const authType: SnowflakeAuthType = selectedAuthType ?? defaultAuthType
+  const authType = selectedAuthType ?? defaultAuthType
   const connectionSubmit = useSavedConnectionSubmit({
     providerId: SNOWFLAKE_SERVICE_ACCOUNT_PROVIDER_ID,
     workspaceId,
@@ -178,6 +179,7 @@ function SnowflakeAccountForm({
     connectionSubmit.isPending || createCredential.isPending || updateCredential.isPending
   const disabled =
     isPending ||
+    !authType ||
     !normalizedHost ||
     (authType === 'programmatic_access_token' && !token.trim()) ||
     (authType === 'key_pair' && (!user.trim() || !privateKey.trim())) ||
@@ -277,15 +279,15 @@ function SnowflakeAccountForm({
         <ChipModalField
           type='dropdown'
           title='Authentication'
-          value={authType}
+          value={authType ?? undefined}
           onChange={(value) => {
             setSelectedAuthType(value as SnowflakeAuthType)
             clearErrors()
           }}
           options={authOptions}
           placeholder='Choose how Sim signs in to Snowflake'
           required
-          hint={AUTH_TYPE_HINTS[authType]}
+          hint={authType ? AUTH_TYPE_HINTS[authType] : undefined}
         />
         <ChipModalField
           type='input'

apps/sim/lib/credentials/snowflake-service-account.server.test.ts

@@ -255,7 +255,6 @@ describe('a Snowflake key-pair connection', () => {
   const FINGERPRINT = `SHA256:${createHash('sha256')
     .update(publicKey.export({ type: 'spki', format: 'der' }))
     .digest('base64')}`
-  const fetchMock = vi.fn<typeof fetch>()
 
   function keyPairConfig(fields: Record<string, unknown> = {}) {
     return JSON.stringify({
@@ -267,34 +266,31 @@ describe('a Snowflake key-pair connection', () => {
     })
   }
 
-  /** Snowflake's side of key-pair auth: a KEYPAIR_JWT bearer token whose claims and key match. */
+  /** Snowflake's side of key-pair auth, behind the egress guard: a KEYPAIR_JWT bearer token whose claims and key match. */
   function snowflakeAccountTrustingKey(account: string) {
-    fetchMock.mockImplementation(async (_input, init) => {
-      const headers = new Headers(init?.headers)
-      const authorization = headers.get('authorization') ?? ''
-      if (headers.get('x-snowflake-authorization-token-type') !== 'KEYPAIR_JWT') {
-        return jsonResponse({ message: 'Invalid token type' }, 401)
+    secureFetch.mockImplementation(
+      async (_url: string, init: { headers: Record<string, string> }) => {
+        const headers = new Headers(init?.headers)
+        const authorization = headers.get('authorization') ?? ''
+        if (headers.get('x-snowflake-authorization-token-type') !== 'KEYPAIR_JWT') {
+          return jsonResponse({ message: 'Invalid token type' }, 401)
+        }
+        try {
+          const { payload } = await jwtVerify(authorization.replace(/^Bearer /, ''), publicKey, {
+            algorithms: ['RS256'],
+            issuer: `${account}.SIM_WORKFLOWS.${FINGERPRINT}`,
+            subject: `${account}.SIM_WORKFLOWS`,
+            requiredClaims: ['iat', 'exp'],
+          })
+          if ((payload.exp ?? 0) - (payload.iat ?? 0) > 3600) throw new Error('lifetime too long')
+        } catch {
+          return jsonResponse({ code: '390144', message: 'JWT token is invalid' }, 401)
+        }
+        return jsonResponse({ data: [['SIM_WORKFLOWS', account, 'ANALYST']] })
       }
-      try {
-        const { payload } = await jwtVerify(authorization.replace(/^Bearer /, ''), publicKey, {
-          algorithms: ['RS256'],
-          issuer: `${account}.SIM_WORKFLOWS.${FINGERPRINT}`,
-          subject: `${account}.SIM_WORKFLOWS`,
-          requiredClaims: ['iat', 'exp'],
-        })
-        if ((payload.exp ?? 0) - (payload.iat ?? 0) > 3600) throw new Error('lifetime too long')
-      } catch {
-        return jsonResponse({ code: '390144', message: 'JWT token is invalid' }, 401)
-      }
-      return jsonResponse({ data: [['SIM_WORKFLOWS', account, 'ANALYST']] })
-    })
+    )
   }
 
-  beforeEach(() => {
-    fetchMock.mockReset()
-    vi.stubGlobal('fetch', fetchMock)
-  })
-
   it('signs in with a JWT naming the account, the uppercased user, and the key fingerprint', async () => {
     snowflakeAccountTrustingKey('MYORG-MYACCOUNT')
 
@@ -315,7 +311,7 @@ describe('a Snowflake key-pair connection', () => {
         privateKey: PRIVATE_KEY.trim(),
       },
     })
-    expect(new URL(String(fetchMock.mock.calls[0][0])).href).toBe(
+    expect(new URL(String(secureFetch.mock.calls[0][0])).href).toBe(
       `https://${HOST}/api/v2/statements`
     )
   })
@@ -386,7 +382,7 @@ describe('a Snowflake key-pair connection', () => {
     const saving = verifyAndEncryptSnowflakeConnection(keyPairConfig({ privateKey: key }), owner)
     await expect(saving).rejects.toBeInstanceOf(SnowflakeCredentialError)
     await expect(saving).rejects.toThrow(/key/i)
-    expect(fetchMock).not.toHaveBeenCalled()
+    expect(secureFetch).not.toHaveBeenCalled()
   })
 
   it('refuses a host outside Snowflake before signing anything', async () => {
@@ -395,7 +391,7 @@ describe('a Snowflake key-pair connection', () => {
         verifyAndEncryptSnowflakeConnection(keyPairConfig({ host: 'evil.example.com' }), owner)
       )
     ).toMatch(/host/i)
-    expect(fetchMock).not.toHaveBeenCalled()
+    expect(secureFetch).not.toHaveBeenCalled()
   })
 
   it('signs a fresh JWT for each run with the key-pair token type', async () => {

apps/sim/lib/credentials/token-service-accounts/grants/infisical.test.ts

@@ -1,3 +1,4 @@
+import { jsonResponse } from '@sim/testing/helpers/http'
 import {
   inputValidationMock,
   inputValidationMockFns,
@@ -21,13 +22,6 @@ const FIELDS = {
 
 const secureFetch = inputValidationMockFns.mockSecureFetchWithValidation
 
-function login(body: unknown, status = 200) {
-  return new Response(JSON.stringify(body), {
-    status,
-    headers: { 'Content-Type': 'application/json' },
-  })
-}
-
 describe('Infisical Universal Auth login', () => {
   beforeEach(() => {
     secureFetch.mockReset()
@@ -48,7 +42,7 @@ describe('Infisical Universal Auth login', () => {
 
   it('logs in with a JSON body at the instance through the configured-endpoint guard', async () => {
     secureFetch.mockResolvedValueOnce(
-      login({
+      jsonResponse({
         accessToken: 'machine-token',
         expiresIn: 7200,
         accessTokenMaxTTL: 2_592_000,
@@ -84,7 +78,7 @@ describe('Infisical Universal Auth login', () => {
 
   it('serves a long-lived token for at most a day so an early revocation is picked up', async () => {
     secureFetch.mockResolvedValueOnce(
-      login({ accessToken: 'long-token', expiresIn: 2_592_000, tokenType: 'Bearer' })
+      jsonResponse({ accessToken: 'long-token', expiresIn: 2_592_000, tokenType: 'Bearer' })
     )
 
     const minted = await mintInfisicalUniversalAuthToken(FIELDS)
@@ -95,7 +89,7 @@ describe('Infisical Universal Auth login', () => {
 
   it('serves the token of an identity whose access token TTL is 0 (no expiry) for a day', async () => {
     secureFetch.mockResolvedValueOnce(
-      login({ accessToken: 'unbounded-token', expiresIn: 0, tokenType: 'Bearer' })
+      jsonResponse({ accessToken: 'unbounded-token', expiresIn: 0, tokenType: 'Bearer' })
     )
 
     const minted = await mintInfisicalUniversalAuthToken(FIELDS)
@@ -112,7 +106,7 @@ describe('Infisical Universal Auth login', () => {
     ['a fractional lifetime', { accessToken: 'token', expiresIn: 1.5 }],
     ['a non-object body', ['token']],
   ])('treats %s as the provider failing, not the credentials', async (_case, body) => {
-    secureFetch.mockResolvedValueOnce(login(body))
+    secureFetch.mockResolvedValueOnce(jsonResponse(body))
 
     await expect(mintInfisicalUniversalAuthToken(FIELDS)).rejects.toMatchObject({
       code: 'provider_unavailable',
@@ -128,7 +122,7 @@ describe('Infisical Universal Auth login', () => {
     [500, 'provider_unavailable'],
   ])('classifies a %s login refusal as %s', async (status, code) => {
     secureFetch.mockResolvedValueOnce(
-      login({ statusCode: status, message: 'nope', error: 'Unauthorized' }, status)
+      jsonResponse({ statusCode: status, message: 'nope', error: 'Unauthorized' }, status)
     )
 
     await expect(mintInfisicalUniversalAuthToken(FIELDS)).rejects.toMatchObject({ code, status })

apps/sim/lib/credentials/token-service-accounts/grants/servicenow.test.ts

@@ -1,3 +1,4 @@
+import { jsonResponse } from '@sim/testing/helpers/http'
 import {
   inputValidationMock,
   inputValidationMockFns,
@@ -21,13 +22,6 @@ const FIELDS = {
 
 const secureFetch = inputValidationMockFns.mockSecureFetchWithValidation
 
-function tokenResponse(body: unknown, status = 200) {
-  return new Response(JSON.stringify(body), {
-    status,
-    headers: { 'Content-Type': 'application/json' },
-  })
-}
-
 describe('ServiceNow client credentials grant', () => {
   beforeEach(() => {
     secureFetch.mockReset()
@@ -48,7 +42,7 @@ describe('ServiceNow client credentials grant', () => {
 
   it('posts the grant as a form to the instance oauth_token.do through the egress guard', async () => {
     secureFetch.mockResolvedValueOnce(
-      tokenResponse({ access_token: 'sn-token', expires_in: 1799, token_type: 'Bearer' })
+      jsonResponse({ access_token: 'sn-token', expires_in: 1799, token_type: 'Bearer' })
     )
 
     const minted = await mintServiceNowClientCredentialsToken(FIELDS)

apps/sim/lib/credentials/token-service-accounts/validators/snowflake.test.ts

@@ -10,6 +10,7 @@ vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock
 import { EgressRefusedError } from '@/lib/core/network/transport.server'
 import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
 import {
+  validateSnowflakeKeyPair,
   validateSnowflakeServiceAccount,
   validateSnowflakeWorkloadIdentity,
 } from '@/lib/credentials/token-service-accounts/validators/snowflake'
@@ -117,3 +118,26 @@ describe('validateSnowflakeWorkloadIdentity', () => {
     )
   })
 })
+
+describe('validateSnowflakeKeyPair', () => {
+  beforeEach(() => {
+    vi.stubGlobal('fetch', mockFetch)
+  })
+
+  afterEach(() => {
+    vi.resetAllMocks()
+  })
+
+  it('sends the signed JWT only through the egress guard', async () => {
+    inputValidationMockFns.mockSecureFetchWithValidation.mockRejectedValueOnce(
+      new EgressRefusedError('private address')
+    )
+
+    await expectCode(
+      validateSnowflakeKeyPair({ domain: fields.domain, jwt: 'signed-jwt' }),
+      'host_not_allowed',
+      400
+    )
+    expect(mockFetch).not.toHaveBeenCalled()
+  })
+})