test: add tests for k256 compat, version API, and algorithm validation

kvinwang · kvinwang · commit 762a3a1aa29f · 2026-02-11T03:10:30.000Z
- guest-agent: unit tests for normalize_algorithm, GetKey with k256
  alias, secp256k1_prehashed rejection, ed25519, default algorithm,
  unsupported algorithm, version RPC, and sign with k256 alias
- sdk/rust: integration tests for version(), k256 sign alias
- sdk/go: tests for GetVersion, GetKey k256 alias, unsupported
  algorithm, secp256k1_prehashed rejection, ed25519 validation
- sdk/python: sync/async tests for version(), k256 alias,
  secp256k1_prehashed rejection
- sdk/js: tests for version(), k256 alias, secp256k1_prehashed
  rejection
diff --git a/guest-agent/src/rpc_service.rs b/guest-agent/src/rpc_service.rs
@@ -1155,4 +1155,154 @@ pNs85uhOZE8z2jr8Pg==
         assert!(result.is_err());
         assert_eq!(result.unwrap_err().to_string(), "Unsupported algorithm");
     }
+
+    #[test]
+    fn test_normalize_algorithm() {
+        assert_eq!(normalize_algorithm("k256"), "secp256k1");
+        assert_eq!(normalize_algorithm("secp256k1"), "secp256k1");
+        assert_eq!(normalize_algorithm("ed25519"), "ed25519");
+        assert_eq!(normalize_algorithm(""), "");
+        assert_eq!(normalize_algorithm("unknown"), "unknown");
+    }
+
+    #[tokio::test]
+    async fn test_get_key_k256_alias() {
+        let (state, _guard) = setup_test_state().await;
+        let handler_k256 = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let handler_secp = InternalRpcHandler {
+            state: state.clone(),
+        };
+
+        let req_k256 = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "k256".to_string(),
+        };
+        let req_secp = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "secp256k1".to_string(),
+        };
+
+        let resp_k256 = handler_k256.get_key(req_k256).await.unwrap();
+        let resp_secp = handler_secp.get_key(req_secp).await.unwrap();
+
+        // k256 alias should produce the same key as secp256k1
+        assert_eq!(resp_k256.key, resp_secp.key);
+    }
+
+    #[tokio::test]
+    async fn test_get_key_secp256k1_prehashed_rejected() {
+        let (state, _guard) = setup_test_state().await;
+        let handler = InternalRpcHandler { state };
+
+        let request = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "secp256k1_prehashed".to_string(),
+        };
+
+        let result = handler.get_key(request).await;
+        assert!(result.is_err());
+        assert_eq!(result.unwrap_err().to_string(), "Unsupported algorithm");
+    }
+
+    #[tokio::test]
+    async fn test_get_key_ed25519_success() {
+        let (state, _guard) = setup_test_state().await;
+        let handler = InternalRpcHandler { state };
+
+        let request = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "ed25519".to_string(),
+        };
+
+        let response = handler.get_key(request).await.unwrap();
+        assert!(!response.key.is_empty());
+        assert_eq!(response.signature_chain.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_get_key_default_algorithm() {
+        let (state, _guard) = setup_test_state().await;
+        let handler_default = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let handler_secp = InternalRpcHandler {
+            state: state.clone(),
+        };
+
+        let req_default = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "".to_string(),
+        };
+        let req_secp = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "secp256k1".to_string(),
+        };
+
+        let resp_default = handler_default.get_key(req_default).await.unwrap();
+        let resp_secp = handler_secp.get_key(req_secp).await.unwrap();
+
+        // Empty algorithm should default to secp256k1
+        assert_eq!(resp_default.key, resp_secp.key);
+    }
+
+    #[tokio::test]
+    async fn test_get_key_unsupported_algorithm_fails() {
+        let (state, _guard) = setup_test_state().await;
+        let handler = InternalRpcHandler { state };
+
+        let request = GetKeyArgs {
+            path: "test".to_string(),
+            purpose: "signing".to_string(),
+            algorithm: "rsa".to_string(),
+        };
+
+        let result = handler.get_key(request).await;
+        assert!(result.is_err());
+        assert_eq!(result.unwrap_err().to_string(), "Unsupported algorithm");
+    }
+
+    #[tokio::test]
+    async fn test_version() {
+        let (state, _guard) = setup_test_state().await;
+        let handler = InternalRpcHandler { state };
+
+        let response = handler.version().await.unwrap();
+        assert!(!response.version.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_sign_k256_alias() {
+        let (state, _guard) = setup_test_state().await;
+        let handler_k256 = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let handler_secp = InternalRpcHandler {
+            state: state.clone(),
+        };
+
+        let data = b"test message".to_vec();
+
+        let req_k256 = SignRequest {
+            algorithm: "k256".to_string(),
+            data: data.clone(),
+        };
+        let req_secp = SignRequest {
+            algorithm: "secp256k1".to_string(),
+            data: data.clone(),
+        };
+
+        let resp_k256 = handler_k256.sign(req_k256).await.unwrap();
+        let resp_secp = handler_secp.sign(req_secp).await.unwrap();
+
+        // k256 alias should produce the same public key as secp256k1
+        assert_eq!(resp_k256.public_key, resp_secp.public_key);
+    }
 }
diff --git a/sdk/go/dstack/client_test.go b/sdk/go/dstack/client_test.go
@@ -737,3 +737,63 @@ func TestSignAndVerifySecp256k1Prehashed(t *testing.T) {
 		t.Errorf("expected error to mention '32-byte digest', got: %v", err)
 	}
 }
+
+func TestGetVersion(t *testing.T) {
+	client := dstack.NewDstackClient()
+	resp, err := client.GetVersion(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp.Version == "" {
+		t.Error("expected version to not be empty")
+	}
+}
+
+func TestGetKeyK256Alias(t *testing.T) {
+	client := dstack.NewDstackClient()
+
+	respK256, err := client.GetKey(context.Background(), "/test", "purpose", "k256")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	respSecp, err := client.GetKey(context.Background(), "/test", "purpose", "secp256k1")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// k256 is an alias for secp256k1, should produce the same key
+	if respK256.Key != respSecp.Key {
+		t.Error("expected k256 and secp256k1 to produce the same key")
+	}
+}
+
+func TestGetKeyUnsupportedAlgorithm(t *testing.T) {
+	client := dstack.NewDstackClient()
+	_, err := client.GetKey(context.Background(), "/test", "purpose", "rsa")
+	if err == nil {
+		t.Fatal("expected error for unsupported algorithm")
+	}
+}
+
+func TestGetKeySecp256k1PrehashedRejected(t *testing.T) {
+	client := dstack.NewDstackClient()
+	_, err := client.GetKey(context.Background(), "/test", "purpose", "secp256k1_prehashed")
+	if err == nil {
+		t.Fatal("expected error for secp256k1_prehashed in GetKey")
+	}
+}
+
+func TestGetKeyAlgorithmValidation(t *testing.T) {
+	client := dstack.NewDstackClient()
+
+	// ed25519 should succeed (Version RPC is available on the simulator)
+	resp, err := client.GetKey(context.Background(), "/test", "purpose", "ed25519")
+	if err != nil {
+		t.Fatalf("expected ed25519 to succeed: %v", err)
+	}
+	if resp.Key == "" {
+		t.Error("expected key to not be empty")
+	}
+}
diff --git a/sdk/js/src/__tests__/index.test.ts b/sdk/js/src/__tests__/index.test.ts
@@ -258,6 +258,25 @@ describe('DstackClient', () => {
     })
   })
 
+  it('should be able to get version', async () => {
+    const client = new DstackClient()
+    const result = await client.version()
+    expect(result).toHaveProperty('version')
+    expect(result.version).not.toBe('')
+  })
+
+  it('should get key with k256 alias producing same result as secp256k1', async () => {
+    const client = new DstackClient()
+    const resultK256 = await client.getKey('/test', 'purpose', 'k256')
+    const resultSecp = await client.getKey('/test', 'purpose', 'secp256k1')
+    expect(resultK256.key).toEqual(resultSecp.key)
+  })
+
+  it('should reject secp256k1_prehashed in getKey', async () => {
+    const client = new DstackClient()
+    await expect(() => client.getKey('/test', 'purpose', 'secp256k1_prehashed')).rejects.toThrow()
+  })
+
   describe('deprecated methods with TappdClient', () => {
     it('should support deprecated deriveKey method with warning', async () => {
       const client = new TappdClient()
diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py
@@ -18,6 +18,7 @@
 from dstack_sdk import SignResponse
 from dstack_sdk import TappdClient
 from dstack_sdk import VerifyResponse
+from dstack_sdk import VersionResponse
 from dstack_sdk.dstack_client import InfoResponse
 from dstack_sdk.dstack_client import TcbInfo
 
@@ -537,6 +538,50 @@ async def test_async_tappd_client_tdx_quote_deprecated():
         assert any("tdx_quote is deprecated" in msg for msg in warning_messages)
 
 
+def test_sync_client_version():
+    client = DstackClient()
+    result = client.version()
+    assert isinstance(result, VersionResponse)
+    assert result.version != ""
+
+
+@pytest.mark.asyncio
+async def test_async_client_version():
+    client = AsyncDstackClient()
+    result = await client.version()
+    assert isinstance(result, VersionResponse)
+    assert result.version != ""
+
+
+def test_sync_client_get_key_k256_alias():
+    client = DstackClient()
+    result_k256 = client.get_key(path="/test", purpose="p", algorithm="k256")
+    result_secp = client.get_key(path="/test", purpose="p", algorithm="secp256k1")
+    # k256 is an alias for secp256k1, should produce the same key
+    assert result_k256.decode_key() == result_secp.decode_key()
+
+
+@pytest.mark.asyncio
+async def test_async_client_get_key_k256_alias():
+    client = AsyncDstackClient()
+    result_k256 = await client.get_key(path="/test", purpose="p", algorithm="k256")
+    result_secp = await client.get_key(path="/test", purpose="p", algorithm="secp256k1")
+    assert result_k256.decode_key() == result_secp.decode_key()
+
+
+def test_sync_client_get_key_secp256k1_prehashed_rejected():
+    client = DstackClient()
+    with pytest.raises(Exception):
+        client.get_key(algorithm="secp256k1_prehashed")
+
+
+@pytest.mark.asyncio
+async def test_async_client_get_key_secp256k1_prehashed_rejected():
+    client = AsyncDstackClient()
+    with pytest.raises(Exception):
+        await client.get_key(algorithm="secp256k1_prehashed")
+
+
 @pytest.mark.asyncio
 async def test_async_tappd_client_is_reachable():
     """Test that AsyncTappdClient can check if service is reachable."""
diff --git a/sdk/rust/tests/test_client.rs b/sdk/rust/tests/test_client.rs
@@ -175,3 +175,34 @@ async fn test_async_client_sign_and_verify_secp256k1_prehashed() {
         .unwrap();
     assert!(verify_resp.valid);
 }
+
+#[tokio::test]
+async fn test_async_client_version() {
+    let client = AsyncDstackClient::new(None);
+    let result = client.version().await.unwrap();
+    assert!(!result.version.is_empty());
+}
+
+#[tokio::test]
+async fn test_async_client_get_key_k256_alias() {
+    let client = AsyncDstackClient::new(None);
+    // k256 should work as an alias for secp256k1
+    let result = client.get_key(Some("test".to_string()), None).await.unwrap();
+    assert!(!result.key.is_empty());
+    assert_eq!(result.decode_key().unwrap().len(), 32);
+}
+
+#[tokio::test]
+async fn test_async_client_sign_k256_alias() {
+    let client = AsyncDstackClient::new(None);
+    let data = b"test message".to_vec();
+
+    // Sign with k256 alias
+    let resp_k256 = client.sign("k256", data.clone()).await.unwrap();
+    assert!(!resp_k256.signature.is_empty());
+    assert!(!resp_k256.public_key.is_empty());
+
+    // Sign with secp256k1 should produce the same public key
+    let resp_secp = client.sign("secp256k1", data.clone()).await.unwrap();
+    assert_eq!(resp_k256.public_key, resp_secp.public_key);
+}
