feat(sdk): validate algorithm in getKey by checking OS version

kvinwang · kvinwang · commit 2a8abb1c0088 · 2026-02-11T02:54:44.000Z
When getKey is called with a non-secp256k1 algorithm (e.g. ed25519),
the SDK calls the Version RPC to check if the OS supports it. If the
Version RPC is unavailable (OS &lt;= 0.5.6), the call is rejected with a
clear error instead of silently returning the wrong key type.
diff --git a/sdk/go/dstack/client.go b/sdk/go/dstack/client.go
@@ -369,8 +369,33 @@ func (c *DstackClient) GetTlsKey(
 	return &response, nil
 }
 
+// requiresVersionCheck returns true for algorithms that need OS >= 0.5.7.
+func requiresVersionCheck(algorithm string) bool {
+	switch algorithm {
+	case "secp256k1", "secp256k1_prehashed", "k256", "":
+		return false
+	default:
+		return true
+	}
+}
+
+// ensureAlgorithmSupported checks the OS version when a non-secp256k1 algorithm is requested.
+// On old OS (no Version RPC), it returns an error to prevent silent key type mismatch.
+func (c *DstackClient) ensureAlgorithmSupported(ctx context.Context, algorithm string) error {
+	if !requiresVersionCheck(algorithm) {
+		return nil
+	}
+	if _, err := c.GetVersion(ctx); err != nil {
+		return fmt.Errorf("algorithm %q is not supported: OS version too old (Version RPC unavailable)", algorithm)
+	}
+	return nil
+}
+
 // Gets a key from the dstack service.
 func (c *DstackClient) GetKey(ctx context.Context, path string, purpose string, algorithm string) (*GetKeyResponse, error) {
+	if err := c.ensureAlgorithmSupported(ctx, algorithm); err != nil {
+		return nil, err
+	}
 	payload := map[string]interface{}{
 		"path":      path,
 		"purpose":   purpose,
diff --git a/sdk/js/src/index.ts b/sdk/js/src/index.ts
@@ -177,6 +177,8 @@ export interface TlsKeyOptions {
   usageClientAuth?: boolean;
 }
 
+const SECP256K1_ALGORITHMS = new Set(['secp256k1', 'secp256k1_prehashed', 'k256', ''])
+
 export class DstackClient<T extends TcbInfo = TcbInfoV05x> {
   protected endpoint: string
 
@@ -202,7 +204,17 @@ export class DstackClient<T extends TcbInfo = TcbInfoV05x> {
     this.endpoint = endpoint
   }
 
+  private async ensureAlgorithmSupported(algorithm: string): Promise<void> {
+    if (SECP256K1_ALGORITHMS.has(algorithm)) return
+    try {
+      await this.version()
+    } catch {
+      throw new Error(`algorithm "${algorithm}" is not supported: OS version too old (Version RPC unavailable)`)
+    }
+  }
+
   async getKey(path: string, purpose: string = '', algorithm: string = 'secp256k1'): Promise<GetKeyResponse> {
+    await this.ensureAlgorithmSupported(algorithm)
     const payload = JSON.stringify({
       path: path,
       purpose: purpose,
diff --git a/sdk/python/src/dstack_sdk/dstack_client.py b/sdk/python/src/dstack_sdk/dstack_client.py
@@ -380,13 +380,26 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
                 self._sync_client.close()
                 self._sync_client = None
 
+    async def _ensure_algorithm_supported(self, algorithm: str) -> None:
+        """Check OS version when a non-secp256k1 algorithm is requested."""
+        if algorithm in ("secp256k1", "secp256k1_prehashed", "k256", ""):
+            return
+        try:
+            await self.version()
+        except Exception:
+            raise RuntimeError(
+                f'algorithm "{algorithm}" is not supported: '
+                "OS version too old (Version RPC unavailable)"
+            )
+
     async def get_key(
         self,
         path: str | None = None,
         purpose: str | None = None,
         algorithm: str = "secp256k1",
     ) -> GetKeyResponse:
         """Derive a key from the given path, purpose, and algorithm."""
+        await self._ensure_algorithm_supported(algorithm)
         data: Dict[str, Any] = {
             "path": path or "",
             "purpose": purpose or "",
