Add timestamp to GetAppEnvEncryptPubKey

kms/README.md

@@ -179,6 +179,16 @@ The keys are derived with app id which guarantees apps can not get the keys from
 
 The `GetAppEnvEncryptPubKey` RPC is used by the frontend web page to request the app environment encryption public key when deploying a new app. This key is used to encrypt the app environment variables, which can only be decrypted by the app in TEE.
 
+The response includes:
+- `public_key`: The X25519 public key for encrypting environment variables
+- `timestamp`: Unix timestamp (seconds) when the response was generated
+- `signature`: Legacy signature for backward compatibility
+  - Signs: `Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + public_key)`
+- `signature_v1`: New signature with timestamp to prevent replay attacks
+  - Signs: `Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + timestamp_be_bytes + public_key)`
+
+Clients should prefer verifying `signature_v1` with timestamp to protect against replay attacks. Fall back to `signature` only for backward compatibility with older KMS versions.
+
 ### SignCert
 
 The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this RPC, the KMS node will:

kms/rpc/proto/kms_rpc.proto

@@ -19,7 +19,14 @@ message AppId {
 
 message PublicKeyResponse {
   bytes public_key = 1;
+  // Legacy signature without timestamp (for backward compatibility).
+  // Signs: Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + public_key)
   bytes signature = 2;
+  // Unix timestamp in seconds when the response was generated.
+  uint64 timestamp = 3;
+  // New signature with timestamp to prevent replay attacks.
+  // Signs: Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + timestamp_be_bytes + public_key)
+  bytes signature_v1 = 4;
 }
 
 message AppKeyResponse {

kms/src/crypto.rs

@@ -42,3 +42,21 @@ pub(crate) fn sign_message(
     signature_bytes.push(recid.to_byte());
     Ok(signature_bytes)
 }
+
+/// Sign a message with a timestamp to prevent replay attacks.
+/// The signature covers: prefix + ":" + appid + timestamp_be_bytes + message
+pub(crate) fn sign_message_with_timestamp(
+    key: &SigningKey,
+    prefix: &[u8],
+    appid: &[u8],
+    timestamp: u64,
+    message: &[u8],
+) -> Result<Vec<u8>> {
+    let timestamp_bytes = timestamp.to_be_bytes();
+    let digest =
+        Keccak256::new_with_prefix([prefix, b":", appid, &timestamp_bytes[..], message].concat());
+    let (signature, recid) = key.sign_digest_recoverable(digest)?;
+    let mut signature_bytes = signature.to_vec();
+    signature_bytes.push(recid.to_byte());
+    Ok(signature_bytes)
+}

kms/src/main_service.rs

@@ -27,7 +27,7 @@ use upgrade_authority::BootInfo;
 
 use crate::{
     config::KmsConfig,
-    crypto::{derive_k256_key, sign_message},
+    crypto::{derive_k256_key, sign_message, sign_message_with_timestamp},
 };
 
 mod upgrade_authority;
@@ -288,6 +288,12 @@ impl KmsRpc for RpcHandler {
         let pubkey = x25519_dalek::PublicKey::from(&secret);
 
         let public_key = pubkey.to_bytes().to_vec();
+        let timestamp = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .context("System time before UNIX epoch")?
+            .as_secs();
+
+        // Legacy signature (without timestamp) for backward compatibility
         let signature = sign_message(
             &self.state.k256_key,
             b"dstack-env-encrypt-pubkey",
@@ -296,9 +302,21 @@ impl KmsRpc for RpcHandler {
         )
         .context("Failed to sign the public key")?;
 
+        // New signature with timestamp to prevent replay attacks
+        let signature_v1 = sign_message_with_timestamp(
+            &self.state.k256_key,
+            b"dstack-env-encrypt-pubkey",
+            &request.app_id,
+            timestamp,
+            &public_key,
+        )
+        .context("Failed to sign the public key with timestamp")?;
+
         Ok(PublicKeyResponse {
             public_key,
             signature,
+            timestamp,
+            signature_v1,
         })
     }
 

sdk/js/README.md

@@ -219,13 +219,20 @@ These utilities are for deployment scripts, not runtime SDK operations.
 Encrypt secrets before deploying to dstack:
 
 ```typescript
-import { encryptEnvVars, verifyEnvEncryptPublicKey, type EnvVar } from '@phala/dstack-sdk';
+import { encryptEnvVars, verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy, type EnvVar } from '@phala/dstack-sdk';
 
 // Get and verify the KMS public key
-// (obtain public_key and signature from KMS API)
-const kmsIdentity = verifyEnvEncryptPublicKey(publicKeyBytes, signatureBytes, appId);
+// (obtain public_key, signature_v1, and timestamp from KMS API)
+
+// Prefer signature_v1 with timestamp (prevents replay attacks)
+const kmsIdentity = verifyEnvEncryptPublicKey(publicKeyBytes, signatureV1Bytes, appId, timestamp);
 if (!kmsIdentity) {
-  throw new Error('Invalid KMS key');
+  // Fall back to legacy signature for backward compatibility with older KMS
+  const legacyIdentity = verifyEnvEncryptPublicKeyLegacy(publicKeyBytes, signatureBytes, appId);
+  if (!legacyIdentity) {
+    throw new Error('Invalid KMS key');
+  }
+  console.warn('Using legacy signature without timestamp protection');
 }
 
 // Encrypt variables

sdk/js/src/__tests__/browser-compatibility.test.ts

@@ -168,13 +168,14 @@ describe('Browser Compatibility Tests', () => {
     const testPublicKey = new Uint8Array(32).fill(1) // 32 bytes
     const testSignature = new Uint8Array(65).fill(2) // 65 bytes
     const testAppId = 'test-app-id'
+    const testTimestamp = BigInt(Math.floor(Date.now() / 1000))
 
     it('should accept the same input parameters', async () => {
-      const nodeResult = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, testSignature, testAppId
+      const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
+        testPublicKey, testSignature, testAppId, testTimestamp
       )
       const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, testSignature, testAppId
+        testPublicKey, testSignature, testAppId, testTimestamp
       )
 
       // Both should return string or null
@@ -187,18 +188,18 @@ describe('Browser Compatibility Tests', () => {
       const invalidSignature = new Uint8Array(32) // Wrong size
 
       // Both should handle invalid inputs similarly
-      const nodeResult1 = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        invalidPublicKey, testSignature, testAppId
+      const nodeResult1 = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
+        invalidPublicKey, testSignature, testAppId, testTimestamp
       )
       const browserResult1 = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        invalidPublicKey, testSignature, testAppId
+        invalidPublicKey, testSignature, testAppId, testTimestamp
       )
 
-      const nodeResult2 = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, invalidSignature, testAppId
+      const nodeResult2 = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
+        testPublicKey, invalidSignature, testAppId, testTimestamp
       )
       const browserResult2 = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, invalidSignature, testAppId
+        testPublicKey, invalidSignature, testAppId, testTimestamp
       )
 
       // Both should return null for invalid inputs (or handle errors consistently)
@@ -209,28 +210,45 @@ describe('Browser Compatibility Tests', () => {
     })
 
     it('should handle empty/invalid app ID consistently', async () => {
-      const nodeResult = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, testSignature, ''
+      const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
+        testPublicKey, testSignature, '', testTimestamp
       )
       const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
-        testPublicKey, testSignature, ''
+        testPublicKey, testSignature, '', testTimestamp
       )
 
       expect(nodeResult).toBeNull()
       expect(browserResult).toBeNull()
     })
+
+    it('should have matching legacy function exports', async () => {
+      // Test legacy functions exist and have the same interface
+      expect(typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy).toBe('function')
+      expect(typeof browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy).toBe('function')
+
+      const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy(
+        testPublicKey, testSignature, testAppId
+      )
+      const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy(
+        testPublicKey, testSignature, testAppId
+      )
+
+      expect(nodeResult === null || typeof nodeResult === 'string').toBeTruthy()
+      expect(browserResult === null || typeof browserResult === 'string').toBeTruthy()
+    })
   })
 
   describe('Function Signatures', () => {
     it('should have matching function signatures', () => {
       // These checks ensure TypeScript compatibility
       const nodeEncryptFn: typeof nodeEncryptEnvVars.encryptEnvVars = browserEncryptEnvVars.encryptEnvVars
       const nodeHashFn: typeof nodeGetComposeHash.getComposeHash = browserGetComposeHash.getComposeHash
-      const nodeVerifyFn: typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey = browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey
+      // Note: verify functions have slightly different signatures (sync vs async) but same parameters
+      expect(typeof browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey).toBe('function')
+      expect(typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey).toBe('function')
 
       expect(typeof nodeEncryptFn).toBe('function')
       expect(typeof nodeHashFn).toBe('function')
-      expect(typeof nodeVerifyFn).toBe('function')
     })
   })
 })

sdk/js/src/__tests__/verify-env-encrypt-public-key.test.ts

@@ -2,47 +2,136 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-import { verifyEnvEncryptPublicKey } from '../verify-env-encrypt-public-key'
-import { describe, it, expect } from 'vitest'
+import { verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy } from '../verify-env-encrypt-public-key'
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
 
-describe('verifySignature', () => {
+describe('verifyEnvEncryptPublicKeyLegacy', () => {
   it('should verify signature correctly with example data', () => {
     const publicKey = new Uint8Array(Buffer.from('e33a1832c6562067ff8f844a61e51ad051f1180b66ec2551fb0251735f3ee90a', 'hex'))
     const signature = new Uint8Array(Buffer.from('8542c49081fbf4e03f62034f13fbf70630bdf256a53032e38465a27c36fd6bed7a5e7111652004aef37f7fd92fbfc1285212c4ae6a6154203a48f5e16cad2cef00', 'hex'))
     const appId = '00'.repeat(20)
-    
-    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
-    
+
+    const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
+
     expect(result).toBe('0x0217610d74cbd39b6143842c6d8bc310d79da1d82cc9d17f8876376221eda0c38f')
   })
 
   it('should handle 0x prefix in app_id', () => {
     const publicKey = new Uint8Array(Buffer.from('e33a1832c6562067ff8f844a61e51ad051f1180b66ec2551fb0251735f3ee90a', 'hex'))
     const signature = new Uint8Array(Buffer.from('8542c49081fbf4e03f62034f13fbf70630bdf256a53032e38465a27c36fd6bed7a5e7111652004aef37f7fd92fbfc1285212c4ae6a6154203a48f5e16cad2cef00', 'hex'))
     const appId = '0x' + '00'.repeat(20)
-    
-    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
-    
+
+    const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
+
     expect(result).toBe('0x0217610d74cbd39b6143842c6d8bc310d79da1d82cc9d17f8876376221eda0c38f')
   })
 
   it('should return null for invalid signature length', () => {
     const publicKey = new Uint8Array(32)
     const signature = new Uint8Array(64) // Wrong length
     const appId = '00'.repeat(20)
-    
-    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
-    
+
+    const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
+
     expect(result).toBeNull()
   })
 
   it('should return null for invalid signature data', () => {
     const publicKey = new Uint8Array(32)
     const signature = new Uint8Array(65) // All zeros
     const appId = '00'.repeat(20)
-
-    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
-
+
+    const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
+
+    expect(result).toBeNull()
+  })
+})
+
+describe('verifyEnvEncryptPublicKey with timestamp', () => {
+  beforeEach(() => {
+    // Mock Date.now to return a fixed timestamp for testing
+    vi.useFakeTimers()
+    vi.setSystemTime(new Date('2024-01-15T12:00:00Z'))
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it('should return null for invalid signature length', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(64) // Wrong length
+    const appId = '00'.repeat(20)
+    const timestamp = BigInt(Math.floor(Date.now() / 1000))
+
+    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
+
+    expect(result).toBeNull()
+  })
+
+  it('should return null for stale timestamp', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(65)
+    const appId = '00'.repeat(20)
+    // Timestamp from 10 minutes ago (600 seconds)
+    const timestamp = BigInt(Math.floor(Date.now() / 1000)) - 600n
+
+    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
+
+    expect(result).toBeNull()
+  })
+
+  it('should return null for timestamp too far in the future', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(65)
+    const appId = '00'.repeat(20)
+    // Timestamp 2 minutes in the future
+    const timestamp = BigInt(Math.floor(Date.now() / 1000)) + 120n
+
+    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
+
+    expect(result).toBeNull()
+  })
+
+  it('should accept timestamp within allowed clock skew (future)', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(65) // Invalid signature, but we're testing timestamp check
+    const appId = '00'.repeat(20)
+    // Timestamp 30 seconds in the future (within 60s skew)
+    const timestamp = BigInt(Math.floor(Date.now() / 1000)) + 30n
+
+    // Will return null due to invalid signature, not timestamp
+    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
+
+    // The function should not reject due to timestamp, but will fail signature verification
+    expect(result).toBeNull()
+  })
+
+  it('should accept custom maxAgeSeconds option', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(65)
+    const appId = '00'.repeat(20)
+    // Timestamp from 400 seconds ago (would fail default 300s, but pass 600s)
+    const timestamp = BigInt(Math.floor(Date.now() / 1000)) - 400n
+
+    // With default maxAge (300s), this should fail due to stale timestamp
+    const result1 = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
+    expect(result1).toBeNull()
+
+    // With extended maxAge (600s), it would pass timestamp check but fail signature
+    const result2 = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp, { maxAgeSeconds: 600 })
+    // Still null due to invalid signature data, but the timestamp check passed
+    expect(result2).toBeNull()
+  })
+
+  it('should accept number timestamp', () => {
+    const publicKey = new Uint8Array(32)
+    const signature = new Uint8Array(65)
+    const appId = '00'.repeat(20)
+    const timestamp = Math.floor(Date.now() / 1000) // number instead of bigint
+
+    // Will return null due to invalid signature, but should handle number timestamp
+    const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
     expect(result).toBeNull()
   })
-}) 
+})

sdk/js/src/index.ts

@@ -6,7 +6,8 @@ import fs from 'fs'
 import crypto from 'crypto'
 import { send_rpc_request } from './send-rpc-request'
 export { getComposeHash } from './get-compose-hash'
-export { verifyEnvEncryptPublicKey } from './verify-env-encrypt-public-key'
+export { verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy } from './verify-env-encrypt-public-key'
+export type { VerifyOptions } from './verify-env-encrypt-public-key'
 
 export interface GetTlsKeyResponse {
   __name__: Readonly<'GetTlsKeyResponse'>