Merge pull request #958 from Dstack-TEE/codex/feat-kms-finalized-ethereum-auth

[STACKED on #957] feat(kms): authorize from finalized Ethereum snapshots

Author: kvinwang

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' });
+  });
+});

dstack/kms/auth-eth-bun/index.ts

@@ -113,10 +113,31 @@ const DSTACK_KMS_ABI = [
 class EthereumBackend {
   private client: ReturnType<typeof createPublicClient>;
   private kmsContractAddr: Address;
+  private expectedChainId?: number;
+  private finalityConfirmations: bigint;
 
-  constructor(client: ReturnType<typeof createPublicClient>, kmsContractAddr: string) {
+  constructor(
+    client: ReturnType<typeof createPublicClient>,
+    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<bigint> {
+    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,28 +163,32 @@ 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
       });
     }
 
     const [isAllowed, reason] = response;
     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 {