From 1f60b8af97b15cb8e33f68c44f602cf69d79bd7a Mon Sep 17 00:00:00 2001
From: Divy Srivastava <dj.srivastava23@gmail.com>
Date: Thu, 21 Mar 2024 14:11:54 +0530
Subject: [PATCH] fix(ext/node): ECDH.publicKey() point encoding  (#23013)

---
 ext/node/lib.rs                               |  1 +
 ext/node/ops/crypto/mod.rs                    | 69 +++++++++++++++++++
 .../internal/crypto/diffiehellman.ts          | 14 ++--
 tests/unit_node/crypto/crypto_key_test.ts     | 13 ++++
 4 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/ext/node/lib.rs b/ext/node/lib.rs
index 16e69250b7..cf63a57855 100644
--- a/ext/node/lib.rs
+++ b/ext/node/lib.rs
@@ -240,6 +240,7 @@ deno_core::extension!(deno_node,
     ops::crypto::op_node_ecdh_generate_keys,
     ops::crypto::op_node_ecdh_compute_secret,
     ops::crypto::op_node_ecdh_compute_public_key,
+    ops::crypto::op_node_ecdh_encode_pubkey,
     ops::crypto::x509::op_node_x509_parse,
     ops::crypto::x509::op_node_x509_ca,
     ops::crypto::x509::op_node_x509_check_email,
diff --git a/ext/node/ops/crypto/mod.rs b/ext/node/ops/crypto/mod.rs
index 7ea96c0310..ed1b7fc757 100644
--- a/ext/node/ops/crypto/mod.rs
+++ b/ext/node/ops/crypto/mod.rs
@@ -1047,6 +1047,75 @@ pub async fn op_node_scrypt_async(
   .await?
 }
 
+#[op2]
+#[buffer]
+pub fn op_node_ecdh_encode_pubkey(
+  #[string] curve: &str,
+  #[buffer] pubkey: &[u8],
+  compress: bool,
+) -> Result<Vec<u8>, AnyError> {
+  use elliptic_curve::sec1::FromEncodedPoint;
+
+  match curve {
+    "secp256k1" => {
+      let pubkey =
+        elliptic_curve::PublicKey::<k256::Secp256k1>::from_encoded_point(
+          &elliptic_curve::sec1::EncodedPoint::<k256::Secp256k1>::from_bytes(
+            pubkey,
+          )?,
+        );
+      // CtOption does not expose its variants.
+      if pubkey.is_none().into() {
+        return Err(type_error("Invalid public key"));
+      }
+
+      let pubkey = pubkey.unwrap();
+
+      Ok(pubkey.to_encoded_point(compress).as_ref().to_vec())
+    }
+    "prime256v1" | "secp256r1" => {
+      let pubkey = elliptic_curve::PublicKey::<NistP256>::from_encoded_point(
+        &elliptic_curve::sec1::EncodedPoint::<NistP256>::from_bytes(pubkey)?,
+      );
+      // CtOption does not expose its variants.
+      if pubkey.is_none().into() {
+        return Err(type_error("Invalid public key"));
+      }
+
+      let pubkey = pubkey.unwrap();
+
+      Ok(pubkey.to_encoded_point(compress).as_ref().to_vec())
+    }
+    "secp384r1" => {
+      let pubkey = elliptic_curve::PublicKey::<NistP384>::from_encoded_point(
+        &elliptic_curve::sec1::EncodedPoint::<NistP384>::from_bytes(pubkey)?,
+      );
+      // CtOption does not expose its variants.
+      if pubkey.is_none().into() {
+        return Err(type_error("Invalid public key"));
+      }
+
+      let pubkey = pubkey.unwrap();
+
+      Ok(pubkey.to_encoded_point(compress).as_ref().to_vec())
+    }
+    "secp224r1" => {
+      let pubkey = elliptic_curve::PublicKey::<NistP224>::from_encoded_point(
+        &elliptic_curve::sec1::EncodedPoint::<NistP224>::from_bytes(pubkey)?,
+      );
+      // CtOption does not expose its variants.
+      if pubkey.is_none().into() {
+        return Err(type_error("Invalid public key"));
+      }
+
+      let pubkey = pubkey.unwrap();
+
+      Ok(pubkey.to_encoded_point(compress).as_ref().to_vec())
+    }
+    &_ => Err(type_error("Unsupported curve")),
+  }
+}
+
 #[op2(fast)]
 pub fn op_node_ecdh_generate_keys(
   #[string] curve: &str,
diff --git a/ext/node/polyfills/internal/crypto/diffiehellman.ts b/ext/node/polyfills/internal/crypto/diffiehellman.ts
index da79077343..6058433ba9 100644
--- a/ext/node/polyfills/internal/crypto/diffiehellman.ts
+++ b/ext/node/polyfills/internal/crypto/diffiehellman.ts
@@ -9,6 +9,7 @@ import {
   op_node_dh_generate2,
   op_node_ecdh_compute_public_key,
   op_node_ecdh_compute_secret,
+  op_node_ecdh_encode_pubkey,
   op_node_ecdh_generate_keys,
   op_node_gen_prime,
 } from "ext:core/ops";
@@ -1239,7 +1240,7 @@ export class ECDH {
     format: ECDHKeyFormat = "uncompressed",
   ): Buffer | string {
     this.#pubbuf = Buffer.alloc(
-      format.trim() == "compressed"
+      format == "compressed"
         ? this.#curve.publicKeySizeCompressed
         : this.#curve.publicKeySize,
     );
@@ -1269,12 +1270,17 @@ export class ECDH {
   getPublicKey(encoding: BinaryToTextEncoding, format?: ECDHKeyFormat): string;
   getPublicKey(
     encoding?: BinaryToTextEncoding,
-    _format?: ECDHKeyFormat,
+    format: ECDHKeyFormat = "uncompressed",
   ): Buffer | string {
+    const pubbuf = Buffer.from(op_node_ecdh_encode_pubkey(
+      this.#curve.name,
+      this.#pubbuf,
+      format == "compressed",
+    ));
     if (encoding !== undefined) {
-      return this.#pubbuf.toString(encoding);
+      return pubbuf.toString(encoding);
     }
-    return this.#pubbuf;
+    return pubbuf;
   }
 
   setPrivateKey(privateKey: ArrayBufferView): void;
diff --git a/tests/unit_node/crypto/crypto_key_test.ts b/tests/unit_node/crypto/crypto_key_test.ts
index de941da88d..ef154fc35b 100644
--- a/tests/unit_node/crypto/crypto_key_test.ts
+++ b/tests/unit_node/crypto/crypto_key_test.ts
@@ -341,3 +341,16 @@ Deno.test("ECDH generateKeys compressed", function () {
   const uncompressedKey = ecdh.generateKeys("binary");
   assertEquals(uncompressedKey.length, 65);
 });
+
+Deno.test("ECDH getPublicKey compressed", function () {
+  const ecdh = createECDH("secp256k1");
+  for (const format of ["compressed", "uncompressed"] as const) {
+    ecdh.generateKeys("binary", format);
+
+    const compressedKey = ecdh.getPublicKey("binary", "compressed");
+    assertEquals(compressedKey.length, 33);
+
+    const uncompressedKey = ecdh.getPublicKey("binary");
+    assertEquals(uncompressedKey.length, 65);
+  }
+});