diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 9977b9582..fd97ffdb6 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -8,11 +8,13 @@ import openApiSpec from './openapi.json'; // Mock viem const mockReadContract = vi.fn(); const mockGetChainId = vi.fn(); +const mockGetBlockNumber = vi.fn(); vi.mock('viem', () => ({ createPublicClient: vi.fn(() => ({ readContract: mockReadContract, getChainId: mockGetChainId, + getBlockNumber: mockGetBlockNumber, })), http: vi.fn(), getContract: vi.fn(), @@ -26,6 +28,8 @@ beforeAll(async () => { process.env.ETH_RPC_URL = 'http://localhost:8545'; process.env.KMS_CONTRACT_ADDR = '0x1234567890123456789012345678901234567890'; process.env.PORT = '3001'; + process.env.ETH_CHAIN_ID = '1337'; + process.env.ETH_FINALITY_CONFIRMATIONS = '2'; // Import the app after mocking const indexModule = await import('./index.ts'); @@ -35,6 +39,8 @@ beforeAll(async () => { beforeEach(() => { // Reset mocks before each test vi.clearAllMocks(); + mockGetChainId.mockResolvedValue(1337); + mockGetBlockNumber.mockResolvedValue(100n); }); describe('API Compatibility Tests', () => { @@ -317,8 +323,9 @@ describe('API Compatibility Tests', () => { expect(data.isAllowed).toBe(false); expect(data.reason).toBe('authorization backend unavailable'); - // Verify that console.error was called for real errors - expect(consoleSpy).toHaveBeenCalledWith('error in KMS boot auth:', expect.any(Error)); + // Diagnostics identify the failing boundary without retaining backend details. + expect(consoleSpy).toHaveBeenCalledWith('KMS authorization backend failed'); + expect(JSON.stringify(consoleSpy.mock.calls)).not.toContain('real error'); consoleSpy.mockRestore(); }); @@ -403,3 +410,158 @@ describe('Hex Decoding Compatibility', () => { expect(response.status).toBe(200); }); }); + +describe('Authorization freshness and domain binding', () => { + const requestBody = { + mrAggregated: '0x' + '11'.repeat(32), + osImageHash: '0x' + '22'.repeat(32), + appId: '0x' + '33'.repeat(20), + composeHash: '0x' + '44'.repeat(32), + instanceId: '0x' + '55'.repeat(20), + deviceId: '0x' + '66'.repeat(32), + }; + + const postApp = (body = requestBody) => appFetch(new Request( + 'http://localhost:3001/bootAuth/app', + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }, + )); + + it('re-evaluates replayed payloads instead of caching an earlier allow', async () => { + let decisions = 0; + mockReadContract.mockImplementation((params) => { + expect(params.address).toBe('0x1234567890123456789012345678901234567890'); + if (params.functionName === 'isAppAllowed') { + decisions += 1; + return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; + } + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + + const first = await postApp(); + const replay = await postApp(); + + expect(await first.json()).toMatchObject({ isAllowed: true, reason: 'initial allow' }); + expect(await replay.json()).toMatchObject({ isAllowed: false, reason: 'policy changed' }); + expect(decisions).toBe(2); + }); + + it('binds changed measurements and identities into distinct contract arguments', async () => { + const calls: unknown[] = []; + mockReadContract.mockImplementation((params) => { + if (params.functionName === 'isAppAllowed') { + calls.push(params.args[0]); + return [true, 'allowed']; + } + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + + await postApp(); + await postApp({ ...requestBody, composeHash: '0x' + '77'.repeat(32) }); + await postApp({ ...requestBody, appId: '0x' + '88'.repeat(20) }); + + expect(calls).toHaveLength(3); + expect(calls[0]).not.toEqual(calls[1]); + expect(calls[0]).not.toEqual(calls[2]); + }); + + it('fails closed during backend interruption and succeeds after recovery', async () => { + mockReadContract.mockRejectedValueOnce(new Error('backend unavailable')); + const interrupted = await postApp(); + expect(await interrupted.json()).toEqual({ + isAllowed: false, + gatewayAppId: '', + reason: 'authorization backend unavailable', + }); + + mockReadContract.mockImplementation((params) => { + if (params.functionName === 'isAppAllowed') return [true, 'recovered']; + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + const recovered = await postApp(); + expect(await recovered.json()).toMatchObject({ isAllowed: true, reason: 'recovered' }); + }); +}); + + +describe('Ethereum finalized snapshot authorization', () => { + const requestBody = { + mrAggregated: '0x' + '11'.repeat(32), + osImageHash: '0x' + '22'.repeat(32), + appId: '0x' + '33'.repeat(20), + composeHash: '0x' + '44'.repeat(32), + instanceId: '0x' + '55'.repeat(20), + deviceId: '0x' + '66'.repeat(32), + }; + + const authorize = () => appFetch(new Request('http://localhost:3001/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(requestBody), + })); + + it('reads the decision and gateway identity from one confirmation-depth snapshot', async () => { + mockGetBlockNumber.mockResolvedValue(100n); + mockReadContract.mockImplementation((params) => { + expect(params.blockNumber).toBe(98n); + if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + + const response = await authorize(); + expect(await response.json()).toMatchObject({ isAllowed: true, reason: 'finalized allow' }); + expect(mockGetBlockNumber).toHaveBeenCalledTimes(1); + }); + + it('re-evaluates the canonical finalized snapshot after a short reorg', async () => { + mockGetBlockNumber.mockResolvedValueOnce(100n).mockResolvedValueOnce(101n); + let decisions = 0; + const observedBlocks: bigint[] = []; + mockReadContract.mockImplementation((params) => { + if (params.functionName === 'isAppAllowed') { + observedBlocks.push(params.blockNumber); + decisions += 1; + return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; + } + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + + const before = await authorize(); + const after = await authorize(); + expect(await before.json()).toMatchObject({ isAllowed: true }); + expect(await after.json()).toMatchObject({ isAllowed: false, reason: 'new canonical deny' }); + expect(observedBlocks).toEqual([98n, 99n]); + }); + + it.each([ + ['wrong chain', () => mockGetChainId.mockResolvedValue(1)], + ['stale head', () => mockGetBlockNumber.mockResolvedValue(1n)], + ['head timeout', () => mockGetBlockNumber.mockRejectedValue(new Error('timeout'))], + ])('fails closed for %s and recovers without retained decisions', async (_name, inject) => { + inject(); + const failed = await authorize(); + expect(await failed.json()).toEqual({ + isAllowed: false, + gatewayAppId: '', + reason: 'authorization backend unavailable', + }); + + mockGetChainId.mockResolvedValue(1337); + mockGetBlockNumber.mockResolvedValue(102n); + mockReadContract.mockImplementation((params) => { + if (params.functionName === 'isAppAllowed') return [true, 'recovered']; + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + throw new Error(`unexpected function ${params.functionName}`); + }); + const recovered = await authorize(); + expect(await recovered.json()).toMatchObject({ isAllowed: true, reason: 'recovered' }); + }); +}); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 5ba7f2e9e..3ea279327 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -113,10 +113,31 @@ const DSTACK_KMS_ABI = [ class EthereumBackend { private client: ReturnType; private kmsContractAddr: Address; + private expectedChainId?: number; + private finalityConfirmations: bigint; - constructor(client: ReturnType, kmsContractAddr: string) { + constructor( + client: ReturnType, + kmsContractAddr: string, + expectedChainId: number | undefined, + finalityConfirmations: bigint, + ) { this.client = client; this.kmsContractAddr = kmsContractAddr as Address; + this.expectedChainId = expectedChainId; + this.finalityConfirmations = finalityConfirmations; + } + + private async finalizedBlockNumber(): Promise { + const chainId = await this.client.getChainId(); + if (this.expectedChainId !== undefined && chainId !== this.expectedChainId) { + throw new Error('authorization backend chain ID mismatch'); + } + const head = await this.client.getBlockNumber(); + if (head < this.finalityConfirmations) { + throw new Error('authorization backend has not reached configured finality'); + } + return head - this.finalityConfirmations; } private decodeHex(hex: string, sz: number = 32): Hex { @@ -142,20 +163,23 @@ class EthereumBackend { advisoryIds: bootInfo.advisoryIds || [] }; + const blockNumber = await this.finalizedBlockNumber(); let response; if (isKms) { response = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, functionName: 'isKmsAllowed', - args: [bootInfoStruct] + args: [bootInfoStruct], + blockNumber }); } else { response = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, functionName: 'isAppAllowed', - args: [bootInfoStruct] + args: [bootInfoStruct], + blockNumber }); } @@ -163,7 +187,8 @@ class EthereumBackend { const gatewayAppId = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, - functionName: 'gatewayAppId' + functionName: 'gatewayAppId', + blockNumber }); return { @@ -203,10 +228,21 @@ const app = new Hono(); // initialize ethereum backend const rpcUrl = process.env.ETH_RPC_URL || 'http://localhost:8545'; const kmsContractAddr = process.env.KMS_CONTRACT_ADDR || '0x0000000000000000000000000000000000000000'; +const parseNonNegativeInteger = (name: string, value: string | undefined): number | undefined => { + if (value === undefined || value === '') return undefined; + if (!/^(?:0|[1-9][0-9]*)$/.test(value)) throw new Error(`${name} must be a non-negative integer`); + const parsed = Number(value); + if (!Number.isSafeInteger(parsed)) throw new Error(`${name} exceeds the safe integer range`); + return parsed; +}; +const expectedChainId = parseNonNegativeInteger('ETH_CHAIN_ID', process.env.ETH_CHAIN_ID); +const finalityConfirmations = BigInt( + parseNonNegativeInteger('ETH_FINALITY_CONFIRMATIONS', process.env.ETH_FINALITY_CONFIRMATIONS) ?? 0, +); const client = createPublicClient({ transport: http(rpcUrl) }); -const ethereum = new EthereumBackend(client, kmsContractAddr); +const ethereum = new EthereumBackend(client, kmsContractAddr, expectedChainId, finalityConfirmations); const publicRpcEndpoint = (value: string): string => { try {