Merge pull request #957 from Dstack-TEE/codex/fix-kms-node-auth-safety

[STACKED on #956] fix(kms): align Node authorization safety

Author: kvinwang

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

@@ -208,6 +208,7 @@ describe('API Compatibility Tests', () => {
         }));
 
         expect(response.status).toBe(400);
+        expect(await response.json()).toEqual({ isAllowed: false, reason: 'invalid authorization request', gatewayAppId: '' });
       }
       expect(mockReadContract).not.toHaveBeenCalled();
     });

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

@@ -218,6 +218,7 @@ const publicRpcEndpoint = (value: string): string => {
 };
 
 const backendUnavailable = 'authorization backend unavailable';
+const invalidRequest = { isAllowed: false, reason: 'invalid authorization request', gatewayAppId: '' };
 
 // health check and info endpoint
 app.get('/', async (c) => {
@@ -249,7 +250,9 @@ app.get('/', async (c) => {
 
 // app boot authentication
 app.post('/bootAuth/app',
-  zValidator('json', BootInfoSchema),
+  zValidator('json', BootInfoSchema, (result, c) => {
+    if (!result.success) return c.json(invalidRequest, 400);
+  }),
   async (c) => {
     try {
       const bootInfo = c.req.valid('json');
@@ -268,7 +271,9 @@ app.post('/bootAuth/app',
 
 // KMS boot authentication
 app.post('/bootAuth/kms',
-  zValidator('json', BootInfoSchema),
+  zValidator('json', BootInfoSchema, (result, c) => {
+    if (!result.success) return c.json(invalidRequest, 400);
+  }),
   async (c) => {
     try {
       const bootInfo = c.req.valid('json');

dstack/kms/auth-eth/script/Upgrade.s.sol

@@ -58,7 +58,7 @@ contract UpgradeKmsToV2 is Script {
         vm.startBroadcast();
 
         // Upgrade to a specific contract version
-        Upgrades.upgradeProxy(kmsProxy, "contracts/test-utils/DstackKmsV2.sol:DstackKmsV2", "");
+        Upgrades.upgradeProxy(kmsProxy, "DstackKmsV2.sol:DstackKmsV2", "");
 
         vm.stopBroadcast();
 
@@ -76,7 +76,7 @@ contract UpgradeAppToV2 is Script {
         vm.startBroadcast();
 
         // Upgrade to a specific contract version
-        Upgrades.upgradeProxy(appProxy, "contracts/test-utils/DstackAppV2.sol:DstackAppV2", "");
+        Upgrades.upgradeProxy(appProxy, "DstackAppV2.sol:DstackAppV2", "");
 
         vm.stopBroadcast();
 

dstack/kms/auth-eth/src/main.test.ts

@@ -60,6 +60,20 @@ describe('Server', () => {
       expect(mockCheckBoot).toHaveBeenCalledWith(mockBootInfo, false);
     });
 
+    it('should reject oversized and non-hex measurements before backend use', async () => {
+      jest.clearAllMocks();
+      for (const mrAggregated of ['0x' + 'ab'.repeat(33), 'not-hex']) {
+        const response = await app.inject({
+          method: 'POST',
+          url: '/bootAuth/app',
+          payload: { ...mockBootInfo, mrAggregated }
+        });
+        expect(response.statusCode).toBe(400);
+        expect(JSON.parse(response.payload)).toEqual({ isAllowed: false, reason: 'invalid authorization request', gatewayAppId: '' });
+      }
+      expect(app.ethereum.checkBoot).not.toHaveBeenCalled();
+    });
+
     it('should return 400 for invalid boot info', async () => {
       const response = await app.inject({
         method: 'POST',
@@ -111,7 +125,7 @@ describe('Server', () => {
       expect(response.statusCode).toBe(200);
       const result = JSON.parse(response.payload);
       expect(result.isAllowed).toBe(false);
-      expect(result.reason).toMatch(/Test backend error/);
+      expect(result.reason).toBe('authorization backend unavailable');
     });
   });
 });

dstack/kms/auth-eth/src/server.ts

@@ -19,17 +19,26 @@ export async function build(): Promise<FastifyInstance> {
   });
 
   // Register schema for request/response validation
+  const hex = (bytes: number, description: string) => ({
+    type: 'string',
+    pattern: `^(?:0x[0-9a-fA-F]{0,${bytes * 2}}|[0-9a-fA-F]{0,${bytes * 2}})$`,
+    description,
+  });
+
   server.addSchema({
     $id: 'bootInfo',
     type: 'object',
     required: ['mrAggregated', 'osImageHash', 'appId', 'composeHash', 'instanceId', 'deviceId'],
     properties: {
-      mrAggregated: { type: 'string', description: 'Aggregated MR measurement' },
-      osImageHash: { type: 'string', description: 'OS Image hash' },
-      appId: { type: 'string', description: 'Application ID' },
-      composeHash: { type: 'string', description: 'Compose hash' },
-      instanceId: { type: 'string', description: 'Instance ID' },
-      deviceId: { type: 'string', description: 'Device ID' }
+      mrAggregated: hex(32, 'Aggregated MR measurement'),
+      osImageHash: hex(32, 'OS Image hash'),
+      appId: hex(20, 'Application ID'),
+      composeHash: hex(32, 'Compose hash'),
+      instanceId: hex(20, 'Instance ID'),
+      deviceId: hex(32, 'Device ID'),
+      tcbStatus: { type: 'string', maxLength: 128, default: '' },
+      advisoryIds: { type: 'array', maxItems: 128, items: { type: 'string', maxLength: 256 }, default: [] },
+      mrSystem: { ...hex(32, 'System MR measurement'), default: '' },
     }
   });
 
@@ -50,21 +59,45 @@ export async function build(): Promise<FastifyInstance> {
   const provider = new ethers.JsonRpcProvider(rpcUrl);
   server.decorate('ethereum', new EthereumBackend(provider, kmsContractAddr));
 
-  server.get('/', async (request, reply) => {
-    const batch = await Promise.all([
-      server.ethereum.getGatewayAppId(),
-      server.ethereum.getChainId(),
-      server.ethereum.getAppImplementation(),
-    ]);
-    return {
-      status: 'ok',
-      kmsContractAddr: kmsContractAddr,
-      ethRpcUrl: rpcUrl,
-      gatewayAppId: batch[0],
-      chainId: batch[1],
-      appAuthImplementation: batch[2], // NOTE: for backward compatibility
-      appImplementation: batch[2],
-    };
+  const publicRpcEndpoint = (value: string): string => {
+    try {
+      const endpoint = new URL(value);
+      return `${endpoint.protocol}//${endpoint.host}`;
+    } catch {
+      return 'configured';
+    }
+  };
+  const backendUnavailable = 'authorization backend unavailable';
+  const invalidRequest = { isAllowed: false, reason: 'invalid authorization request', gatewayAppId: '' };
+
+  server.setErrorHandler((error, request, reply) => {
+    if (typeof error === 'object' && error !== null && 'validation' in error) {
+      return reply.code(400).send(invalidRequest);
+    }
+    request.log.error('authorization request failed');
+    return reply.code(500).send({ status: 'error', message: backendUnavailable });
+  });
+
+  server.get('/', async (_request, reply) => {
+    try {
+      const batch = await Promise.all([
+        server.ethereum.getGatewayAppId(),
+        server.ethereum.getChainId(),
+        server.ethereum.getAppImplementation(),
+      ]);
+      return {
+        status: 'ok',
+        kmsContractAddr: kmsContractAddr,
+        ethRpcUrl: publicRpcEndpoint(rpcUrl),
+        gatewayAppId: batch[0],
+        chainId: batch[1],
+        appAuthImplementation: batch[2], // NOTE: for backward compatibility
+        appImplementation: batch[2],
+      };
+    } catch {
+      _request.log.error('authorization backend health check failed');
+      return reply.code(500).send({ status: 'error', message: backendUnavailable });
+    }
   });
 
   // Define routes
@@ -82,11 +115,11 @@ export async function build(): Promise<FastifyInstance> {
     try {
       return await server.ethereum.checkBoot(request.body, false);
     } catch (error) {
-      console.error(error);
+      request.log.error('application authorization backend failed');
       reply.code(200).send({
         isAllowed: false,
         gatewayAppId: '',
-        reason: `${error instanceof Error ? error.message : String(error)}`
+        reason: backendUnavailable
       });
     }
   });
@@ -106,12 +139,12 @@ export async function build(): Promise<FastifyInstance> {
       return await server.ethereum.checkBoot(request.body, true);
     } catch (error) {
       if (!(error instanceof Error && "Test backend error" == error.message)) {
-        console.error(error);
+        request.log.error('KMS authorization backend failed');
       }
       reply.code(200).send({
         isAllowed: false,
         gatewayAppId: '',
-        reason: `${error instanceof Error ? error.message : String(error)}`
+        reason: backendUnavailable
       });
     }
   });

dstack/kms/auth-eth/tsconfig.json

@@ -2,17 +2,26 @@
   "compilerOptions": {
     "target": "ES2020",
     "module": "commonjs",
-    "lib": ["es2020"],
+    "lib": [
+      "es2020"
+    ],
     "strict": true,
     "esModuleInterop": true,
     "skipLibCheck": true,
     "forceConsistentCasingInFileNames": true,
     "outDir": "./dist",
-    "rootDir": ".",
+    "rootDir": "./src",
     "resolveJsonModule": true,
-    "types": ["node", "jest"],
+    "types": [
+      "node",
+      "jest"
+    ],
     "baseUrl": "."
   },
-  "include": ["src/**/*"],
-  "exclude": ["node_modules"]
+  "include": [
+    "src/**/*"
+  ],
+  "exclude": [
+    "node_modules"
+  ]
 }

dstack/kms/dstack-app/compose-dev.yaml

@@ -8,17 +8,17 @@ services:
     build:
       context: .
       dockerfile_inline: |
-        FROM node:18-alpine@sha256:06f7bbbcec00dd10c21a3a0962609600159601b5004d84aff142977b449168e9
+        FROM node:20-bookworm-slim@sha256:2cf067cfed83d5ea958367df9f966191a942351a2df77d6f0193e162b5febfc0
         WORKDIR /app
 
-        RUN apk add --no-cache git
+        RUN apt-get update && apt-get install -y --no-install-recommends git ca-certificates && rm -rf /var/lib/apt/lists/*
         RUN git clone ${GIT_REPOSITORY} && \
             cd dstack && \
             git checkout ${GIT_REV}
         WORKDIR /app/dstack/dstack/kms/auth-eth
-        RUN npm install
-        RUN npx tsc --project tsconfig.json
-        CMD node dist/src/main.js
+        RUN npm ci
+        RUN npm run build
+        CMD node dist/main.js
     environment:
       - HOST=0.0.0.0
       - PORT=8000

dstack/kms/dstack-app/docker-compose.yaml

@@ -10,7 +10,7 @@ services:
       dockerfile_inline: |
         FROM rust:1.92.0@sha256:48851a839d6a67370c9dbe0e709bedc138e3e404b161c5233aedcf2b717366e4
         WORKDIR /app
-        RUN apk add --no-cache git build-base openssl-dev protobuf protobuf-dev perl
+        RUN apt-get update && apt-get install -y --no-install-recommends git ca-certificates && rm -rf /var/lib/apt/lists/* build-base openssl-dev protobuf protobuf-dev perl
         RUN git clone https://github.com/a16z/helios && \
             cd helios && \
             git checkout 5c61864a167c16141a9a12b976c0e9398b332f07
@@ -32,18 +32,16 @@ services:
     build:
       context: .
       dockerfile_inline: |
-        FROM node:18-alpine@sha256:06f7bbbcec00dd10c21a3a0962609600159601b5004d84aff142977b449168e9
+        FROM node:20-bookworm-slim@sha256:2cf067cfed83d5ea958367df9f966191a942351a2df77d6f0193e162b5febfc0
         WORKDIR /app
 
-        RUN apk add --no-cache git
+        RUN apt-get update && apt-get install -y --no-install-recommends git ca-certificates && rm -rf /var/lib/apt/lists/*
         RUN git clone https://github.com/Dstack-TEE/dstack.git && \
             cd dstack && \
             git checkout 78057c975fe4b9e21f557fb888d72eeecfb21178
         WORKDIR /app/dstack/kms/auth-eth
-        RUN npm install && \
-            npx hardhat typechain && \
-            npx tsc --project tsconfig.json
-        CMD node dist/src/main.js
+        RUN npm ci && npm run build
+        CMD node dist/main.js
     environment:
       - HOST=0.0.0.0
       - PORT=8000