Skip to content

Commit 76ace38

Browse files
authored
Merge pull request #452 from Dstack-TEE/kms-api-ts
Add timestamp to GetAppEnvEncryptPubKey
2 parents a0e2a82 + 1ca3693 commit 76ace38

14 files changed

Lines changed: 603 additions & 80 deletions

kms/README.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,16 @@ The keys are derived with app id which guarantees apps can not get the keys from
179179

180180
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.
181181

182+
The response includes:
183+
- `public_key`: The X25519 public key for encrypting environment variables
184+
- `timestamp`: Unix timestamp (seconds) when the response was generated
185+
- `signature`: Legacy signature for backward compatibility
186+
- Signs: `Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + public_key)`
187+
- `signature_v1`: New signature with timestamp to prevent replay attacks
188+
- Signs: `Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + timestamp_be_bytes + public_key)`
189+
190+
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.
191+
182192
### SignCert
183193

184194
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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,14 @@ message AppId {
1919

2020
message PublicKeyResponse {
2121
bytes public_key = 1;
22+
// Legacy signature without timestamp (for backward compatibility).
23+
// Signs: Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + public_key)
2224
bytes signature = 2;
25+
// Unix timestamp in seconds when the response was generated.
26+
uint64 timestamp = 3;
27+
// New signature with timestamp to prevent replay attacks.
28+
// Signs: Keccak256("dstack-env-encrypt-pubkey" + ":" + app_id + timestamp_be_bytes + public_key)
29+
bytes signature_v1 = 4;
2330
}
2431

2532
message AppKeyResponse {

kms/src/crypto.rs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,21 @@ pub(crate) fn sign_message(
4242
signature_bytes.push(recid.to_byte());
4343
Ok(signature_bytes)
4444
}
45+
46+
/// Sign a message with a timestamp to prevent replay attacks.
47+
/// The signature covers: prefix + ":" + appid + timestamp_be_bytes + message
48+
pub(crate) fn sign_message_with_timestamp(
49+
key: &SigningKey,
50+
prefix: &[u8],
51+
appid: &[u8],
52+
timestamp: u64,
53+
message: &[u8],
54+
) -> Result<Vec<u8>> {
55+
let timestamp_bytes = timestamp.to_be_bytes();
56+
let digest =
57+
Keccak256::new_with_prefix([prefix, b":", appid, &timestamp_bytes[..], message].concat());
58+
let (signature, recid) = key.sign_digest_recoverable(digest)?;
59+
let mut signature_bytes = signature.to_vec();
60+
signature_bytes.push(recid.to_byte());
61+
Ok(signature_bytes)
62+
}

kms/src/main_service.rs

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use upgrade_authority::BootInfo;
2727

2828
use crate::{
2929
config::KmsConfig,
30-
crypto::{derive_k256_key, sign_message},
30+
crypto::{derive_k256_key, sign_message, sign_message_with_timestamp},
3131
};
3232

3333
mod upgrade_authority;
@@ -288,6 +288,12 @@ impl KmsRpc for RpcHandler {
288288
let pubkey = x25519_dalek::PublicKey::from(&secret);
289289

290290
let public_key = pubkey.to_bytes().to_vec();
291+
let timestamp = std::time::SystemTime::now()
292+
.duration_since(std::time::UNIX_EPOCH)
293+
.context("System time before UNIX epoch")?
294+
.as_secs();
295+
296+
// Legacy signature (without timestamp) for backward compatibility
291297
let signature = sign_message(
292298
&self.state.k256_key,
293299
b"dstack-env-encrypt-pubkey",
@@ -296,9 +302,21 @@ impl KmsRpc for RpcHandler {
296302
)
297303
.context("Failed to sign the public key")?;
298304

305+
// New signature with timestamp to prevent replay attacks
306+
let signature_v1 = sign_message_with_timestamp(
307+
&self.state.k256_key,
308+
b"dstack-env-encrypt-pubkey",
309+
&request.app_id,
310+
timestamp,
311+
&public_key,
312+
)
313+
.context("Failed to sign the public key with timestamp")?;
314+
299315
Ok(PublicKeyResponse {
300316
public_key,
301317
signature,
318+
timestamp,
319+
signature_v1,
302320
})
303321
}
304322

sdk/js/README.md

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -219,13 +219,20 @@ These utilities are for deployment scripts, not runtime SDK operations.
219219
Encrypt secrets before deploying to dstack:
220220

221221
```typescript
222-
import { encryptEnvVars, verifyEnvEncryptPublicKey, type EnvVar } from '@phala/dstack-sdk';
222+
import { encryptEnvVars, verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy, type EnvVar } from '@phala/dstack-sdk';
223223

224224
// Get and verify the KMS public key
225-
// (obtain public_key and signature from KMS API)
226-
const kmsIdentity = verifyEnvEncryptPublicKey(publicKeyBytes, signatureBytes, appId);
225+
// (obtain public_key, signature_v1, and timestamp from KMS API)
226+
227+
// Prefer signature_v1 with timestamp (prevents replay attacks)
228+
const kmsIdentity = verifyEnvEncryptPublicKey(publicKeyBytes, signatureV1Bytes, appId, timestamp);
227229
if (!kmsIdentity) {
228-
throw new Error('Invalid KMS key');
230+
// Fall back to legacy signature for backward compatibility with older KMS
231+
const legacyIdentity = verifyEnvEncryptPublicKeyLegacy(publicKeyBytes, signatureBytes, appId);
232+
if (!legacyIdentity) {
233+
throw new Error('Invalid KMS key');
234+
}
235+
console.warn('Using legacy signature without timestamp protection');
229236
}
230237

231238
// Encrypt variables

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

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -168,13 +168,14 @@ describe('Browser Compatibility Tests', () => {
168168
const testPublicKey = new Uint8Array(32).fill(1) // 32 bytes
169169
const testSignature = new Uint8Array(65).fill(2) // 65 bytes
170170
const testAppId = 'test-app-id'
171+
const testTimestamp = BigInt(Math.floor(Date.now() / 1000))
171172

172173
it('should accept the same input parameters', async () => {
173-
const nodeResult = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
174-
testPublicKey, testSignature, testAppId
174+
const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
175+
testPublicKey, testSignature, testAppId, testTimestamp
175176
)
176177
const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
177-
testPublicKey, testSignature, testAppId
178+
testPublicKey, testSignature, testAppId, testTimestamp
178179
)
179180

180181
// Both should return string or null
@@ -187,18 +188,18 @@ describe('Browser Compatibility Tests', () => {
187188
const invalidSignature = new Uint8Array(32) // Wrong size
188189

189190
// Both should handle invalid inputs similarly
190-
const nodeResult1 = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
191-
invalidPublicKey, testSignature, testAppId
191+
const nodeResult1 = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
192+
invalidPublicKey, testSignature, testAppId, testTimestamp
192193
)
193194
const browserResult1 = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
194-
invalidPublicKey, testSignature, testAppId
195+
invalidPublicKey, testSignature, testAppId, testTimestamp
195196
)
196197

197-
const nodeResult2 = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
198-
testPublicKey, invalidSignature, testAppId
198+
const nodeResult2 = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
199+
testPublicKey, invalidSignature, testAppId, testTimestamp
199200
)
200201
const browserResult2 = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
201-
testPublicKey, invalidSignature, testAppId
202+
testPublicKey, invalidSignature, testAppId, testTimestamp
202203
)
203204

204205
// Both should return null for invalid inputs (or handle errors consistently)
@@ -209,28 +210,45 @@ describe('Browser Compatibility Tests', () => {
209210
})
210211

211212
it('should handle empty/invalid app ID consistently', async () => {
212-
const nodeResult = await nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
213-
testPublicKey, testSignature, ''
213+
const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
214+
testPublicKey, testSignature, '', testTimestamp
214215
)
215216
const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey(
216-
testPublicKey, testSignature, ''
217+
testPublicKey, testSignature, '', testTimestamp
217218
)
218219

219220
expect(nodeResult).toBeNull()
220221
expect(browserResult).toBeNull()
221222
})
223+
224+
it('should have matching legacy function exports', async () => {
225+
// Test legacy functions exist and have the same interface
226+
expect(typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy).toBe('function')
227+
expect(typeof browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy).toBe('function')
228+
229+
const nodeResult = nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy(
230+
testPublicKey, testSignature, testAppId
231+
)
232+
const browserResult = await browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKeyLegacy(
233+
testPublicKey, testSignature, testAppId
234+
)
235+
236+
expect(nodeResult === null || typeof nodeResult === 'string').toBeTruthy()
237+
expect(browserResult === null || typeof browserResult === 'string').toBeTruthy()
238+
})
222239
})
223240

224241
describe('Function Signatures', () => {
225242
it('should have matching function signatures', () => {
226243
// These checks ensure TypeScript compatibility
227244
const nodeEncryptFn: typeof nodeEncryptEnvVars.encryptEnvVars = browserEncryptEnvVars.encryptEnvVars
228245
const nodeHashFn: typeof nodeGetComposeHash.getComposeHash = browserGetComposeHash.getComposeHash
229-
const nodeVerifyFn: typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey = browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey
246+
// Note: verify functions have slightly different signatures (sync vs async) but same parameters
247+
expect(typeof browserVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey).toBe('function')
248+
expect(typeof nodeVerifyEnvEncryptPublicKey.verifyEnvEncryptPublicKey).toBe('function')
230249

231250
expect(typeof nodeEncryptFn).toBe('function')
232251
expect(typeof nodeHashFn).toBe('function')
233-
expect(typeof nodeVerifyFn).toBe('function')
234252
})
235253
})
236254
})

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

Lines changed: 105 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,47 +2,136 @@
22
//
33
// SPDX-License-Identifier: Apache-2.0
44

5-
import { verifyEnvEncryptPublicKey } from '../verify-env-encrypt-public-key'
6-
import { describe, it, expect } from 'vitest'
5+
import { verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy } from '../verify-env-encrypt-public-key'
6+
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
77

8-
describe('verifySignature', () => {
8+
describe('verifyEnvEncryptPublicKeyLegacy', () => {
99
it('should verify signature correctly with example data', () => {
1010
const publicKey = new Uint8Array(Buffer.from('e33a1832c6562067ff8f844a61e51ad051f1180b66ec2551fb0251735f3ee90a', 'hex'))
1111
const signature = new Uint8Array(Buffer.from('8542c49081fbf4e03f62034f13fbf70630bdf256a53032e38465a27c36fd6bed7a5e7111652004aef37f7fd92fbfc1285212c4ae6a6154203a48f5e16cad2cef00', 'hex'))
1212
const appId = '00'.repeat(20)
13-
14-
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
15-
13+
14+
const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
15+
1616
expect(result).toBe('0x0217610d74cbd39b6143842c6d8bc310d79da1d82cc9d17f8876376221eda0c38f')
1717
})
1818

1919
it('should handle 0x prefix in app_id', () => {
2020
const publicKey = new Uint8Array(Buffer.from('e33a1832c6562067ff8f844a61e51ad051f1180b66ec2551fb0251735f3ee90a', 'hex'))
2121
const signature = new Uint8Array(Buffer.from('8542c49081fbf4e03f62034f13fbf70630bdf256a53032e38465a27c36fd6bed7a5e7111652004aef37f7fd92fbfc1285212c4ae6a6154203a48f5e16cad2cef00', 'hex'))
2222
const appId = '0x' + '00'.repeat(20)
23-
24-
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
25-
23+
24+
const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
25+
2626
expect(result).toBe('0x0217610d74cbd39b6143842c6d8bc310d79da1d82cc9d17f8876376221eda0c38f')
2727
})
2828

2929
it('should return null for invalid signature length', () => {
3030
const publicKey = new Uint8Array(32)
3131
const signature = new Uint8Array(64) // Wrong length
3232
const appId = '00'.repeat(20)
33-
34-
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
35-
33+
34+
const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
35+
3636
expect(result).toBeNull()
3737
})
3838

3939
it('should return null for invalid signature data', () => {
4040
const publicKey = new Uint8Array(32)
4141
const signature = new Uint8Array(65) // All zeros
4242
const appId = '00'.repeat(20)
43-
44-
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId)
45-
43+
44+
const result = verifyEnvEncryptPublicKeyLegacy(publicKey, signature, appId)
45+
46+
expect(result).toBeNull()
47+
})
48+
})
49+
50+
describe('verifyEnvEncryptPublicKey with timestamp', () => {
51+
beforeEach(() => {
52+
// Mock Date.now to return a fixed timestamp for testing
53+
vi.useFakeTimers()
54+
vi.setSystemTime(new Date('2024-01-15T12:00:00Z'))
55+
})
56+
57+
afterEach(() => {
58+
vi.useRealTimers()
59+
})
60+
61+
it('should return null for invalid signature length', () => {
62+
const publicKey = new Uint8Array(32)
63+
const signature = new Uint8Array(64) // Wrong length
64+
const appId = '00'.repeat(20)
65+
const timestamp = BigInt(Math.floor(Date.now() / 1000))
66+
67+
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
68+
69+
expect(result).toBeNull()
70+
})
71+
72+
it('should return null for stale timestamp', () => {
73+
const publicKey = new Uint8Array(32)
74+
const signature = new Uint8Array(65)
75+
const appId = '00'.repeat(20)
76+
// Timestamp from 10 minutes ago (600 seconds)
77+
const timestamp = BigInt(Math.floor(Date.now() / 1000)) - 600n
78+
79+
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
80+
81+
expect(result).toBeNull()
82+
})
83+
84+
it('should return null for timestamp too far in the future', () => {
85+
const publicKey = new Uint8Array(32)
86+
const signature = new Uint8Array(65)
87+
const appId = '00'.repeat(20)
88+
// Timestamp 2 minutes in the future
89+
const timestamp = BigInt(Math.floor(Date.now() / 1000)) + 120n
90+
91+
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
92+
93+
expect(result).toBeNull()
94+
})
95+
96+
it('should accept timestamp within allowed clock skew (future)', () => {
97+
const publicKey = new Uint8Array(32)
98+
const signature = new Uint8Array(65) // Invalid signature, but we're testing timestamp check
99+
const appId = '00'.repeat(20)
100+
// Timestamp 30 seconds in the future (within 60s skew)
101+
const timestamp = BigInt(Math.floor(Date.now() / 1000)) + 30n
102+
103+
// Will return null due to invalid signature, not timestamp
104+
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
105+
106+
// The function should not reject due to timestamp, but will fail signature verification
107+
expect(result).toBeNull()
108+
})
109+
110+
it('should accept custom maxAgeSeconds option', () => {
111+
const publicKey = new Uint8Array(32)
112+
const signature = new Uint8Array(65)
113+
const appId = '00'.repeat(20)
114+
// Timestamp from 400 seconds ago (would fail default 300s, but pass 600s)
115+
const timestamp = BigInt(Math.floor(Date.now() / 1000)) - 400n
116+
117+
// With default maxAge (300s), this should fail due to stale timestamp
118+
const result1 = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
119+
expect(result1).toBeNull()
120+
121+
// With extended maxAge (600s), it would pass timestamp check but fail signature
122+
const result2 = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp, { maxAgeSeconds: 600 })
123+
// Still null due to invalid signature data, but the timestamp check passed
124+
expect(result2).toBeNull()
125+
})
126+
127+
it('should accept number timestamp', () => {
128+
const publicKey = new Uint8Array(32)
129+
const signature = new Uint8Array(65)
130+
const appId = '00'.repeat(20)
131+
const timestamp = Math.floor(Date.now() / 1000) // number instead of bigint
132+
133+
// Will return null due to invalid signature, but should handle number timestamp
134+
const result = verifyEnvEncryptPublicKey(publicKey, signature, appId, timestamp)
46135
expect(result).toBeNull()
47136
})
48-
})
137+
})

sdk/js/src/index.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ import fs from 'fs'
66
import crypto from 'crypto'
77
import { send_rpc_request } from './send-rpc-request'
88
export { getComposeHash } from './get-compose-hash'
9-
export { verifyEnvEncryptPublicKey } from './verify-env-encrypt-public-key'
9+
export { verifyEnvEncryptPublicKey, verifyEnvEncryptPublicKeyLegacy } from './verify-env-encrypt-public-key'
10+
export type { VerifyOptions } from './verify-env-encrypt-public-key'
1011

1112
export interface GetTlsKeyResponse {
1213
__name__: Readonly<'GetTlsKeyResponse'>

0 commit comments

Comments
 (0)