Fix kms signing

Author: danimhr

Cargo.lock

@@ -11,6 +11,7 @@ aws-config = { version = "1.8.0", features = ["behavior-version-latest"] }
 aws-sdk-kms = "1.76.0"
 borsh = "0.9.3"
 clap = { version = "4.5.39", features = ["derive", "env"] }
+der = { version = "0.7.10", features = ["oid", "derive"] }
 hex = { version = "0.4.3", features = ["serde"] }
 prost = "0.14.1"
 reqwest = { version = "0.12.19", features = ["json"] }

Cargo.toml

@@ -1,5 +1,4 @@
 use clap::Parser;
-use reqwest::Url;
 
 #[derive(clap::ValueEnum, Clone, Debug)]
 pub enum Mode {
@@ -26,7 +25,7 @@ pub struct RunOptions {
     /// URI for the signer.
     /// https://github.com/wormhole-foundation/wormhole/blob/main/docs/guardian_signer.md
     #[arg(long = "signer-uri", env = "SIGNER_URI")]
-    pub signer_uri: Url,
+    pub signer_uri: String,
 }
 
 #[derive(Parser, Clone, Debug)]

src/config.rs

@@ -8,6 +8,7 @@ use {
     clap::Parser,
     posted_message::PostedMessageUnreliableData,
     prost::Message,
+    reqwest::Url,
     secp256k1::{rand::rngs::OsRng, Secp256k1},
     sequoia_openpgp::armor::{Kind, Writer},
     serde_wormhole::RawMessage,
@@ -172,11 +173,14 @@ async fn run_listener(input: RunListenerInput) -> Result<(), PubsubClientError>
 }
 
 async fn get_signer(run_options: config::RunOptions) -> anyhow::Result<Arc<dyn Signer>> {
-    match run_options.signer_uri.scheme() {
+    let scheme = run_options.signer_uri.split("://").next().unwrap_or("");
+    println!("Using signer URI: {}", scheme);
+    match scheme {
         "file" => {
             let signer = signer::FileSigner::try_new(
                 run_options
                     .signer_uri
+                    .parse::<Url>()?
                     .to_file_path()
                     .map_err(|_| anyhow::anyhow!("Invalid file path in signer URI"))?,
                 run_options.mode,
@@ -187,9 +191,14 @@ async fn get_signer(run_options: config::RunOptions) -> anyhow::Result<Arc<dyn S
             let arn_string = run_options
                 .signer_uri
                 .as_str()
-                .strip_prefix(&format!("{}://", run_options.signer_uri.scheme()))
+                .strip_prefix("amazonkms://")
                 .ok_or_else(|| anyhow::anyhow!("Invalid Amazon KMS ARN in signer URI"))?;
-            let signer = signer::KMSSigner::try_new(arn_string.to_string()).await?;
+            println!("Using Amazon KMS signer with ARN: {}", arn_string);
+            let mut signer = signer::KMSSigner::try_new(arn_string.to_string()).await?;
+            signer
+                .get_and_cache_public_key()
+                .await
+                .map_err(|e| anyhow::anyhow!("Failed to get public key: {}", e))?;
             Ok(Arc::new(signer))
         }
         _ => Err(anyhow::anyhow!("Unsupported signer URI scheme")),

src/main.rs

@@ -1,3 +1,8 @@
+use der::{
+    asn1::{AnyRef, BitStringRef, UintRef},
+    oid::ObjectIdentifier,
+    Decode, Sequence,
+};
 use std::{
     fs,
     io::{Cursor, Read},
@@ -7,7 +12,10 @@ use std::{
 
 use async_trait::async_trait;
 use prost::Message as ProstMessage;
-use secp256k1::{Message, PublicKey, Secp256k1, SecretKey};
+use secp256k1::{
+    ecdsa::{RecoverableSignature, RecoveryId},
+    Message, PublicKey, Secp256k1, SecretKey,
+};
 use sequoia_openpgp::armor::{Kind, Reader, ReaderMode};
 use sha3::{Digest, Keccak256};
 
@@ -109,15 +117,53 @@ impl Signer for FileSigner {
 pub struct KMSSigner {
     client: aws_sdk_kms::Client,
     arn: aws_arn::ResourceName,
+    public_key: Option<(PublicKey, [u8; 20])>,
 }
 
 impl KMSSigner {
     pub async fn try_new(arn_string: String) -> anyhow::Result<Self> {
         let config = aws_config::load_from_env().await;
         let client = aws_sdk_kms::Client::new(&config);
         let arn = aws_arn::ResourceName::from_str(&arn_string)?;
-        Ok(KMSSigner { client, arn })
+        Ok(KMSSigner {
+            client,
+            arn,
+            public_key: None,
+        })
     }
+
+    pub async fn get_and_cache_public_key(&mut self) -> anyhow::Result<()> {
+        let (public_key, pubkey_evm) = self.get_public_key().await?;
+        self.public_key = Some((public_key, pubkey_evm));
+        Ok(())
+    }
+}
+
+/// X.509 `AlgorithmIdentifier` (same as above)
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Sequence)] // NOTE: added `Sequence`
+pub struct AlgorithmIdentifier<'a> {
+    /// This field contains an ASN.1 `OBJECT IDENTIFIER`, a.k.a. OID.
+    pub algorithm: ObjectIdentifier,
+
+    /// This field is `OPTIONAL` and contains the ASN.1 `ANY` type, which
+    /// in this example allows arbitrary algorithm-defined parameters.
+    pub parameters: Option<AnyRef<'a>>,
+}
+
+/// X.509 `SubjectPublicKeyInfo` (SPKI)
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Sequence)]
+pub struct SubjectPublicKeyInfo<'a> {
+    /// X.509 `AlgorithmIdentifier`
+    pub algorithm: AlgorithmIdentifier<'a>,
+
+    /// Public key data
+    pub subject_public_key: BitStringRef<'a>,
+}
+
+#[derive(Sequence)]
+struct EcdsaSignature<'a> {
+    r: UintRef<'a>,
+    s: UintRef<'a>,
 }
 
 #[async_trait]
@@ -128,23 +174,50 @@ impl Signer for KMSSigner {
             .sign()
             .key_id(self.arn.to_string())
             .message(data.to_vec().into())
+            .message_type(aws_sdk_kms::types::MessageType::Digest)
             .signing_algorithm(aws_sdk_kms::types::SigningAlgorithmSpec::EcdsaSha256)
             .send()
             .await
             .map_err(|e| anyhow::anyhow!("Failed to sign data with KMS: {}", e))?;
-        result
+        let kms_signature = result
             .signature
-            .ok_or_else(|| anyhow::anyhow!("KMS did not return a signature"))
-            .and_then(|sig| {
-                let sig = sig.into_inner();
-                let signature: [u8; 65] = sig
-                    .try_into()
-                    .map_err(|e| anyhow::anyhow!("Failed to convert KMS signature: {:?}", e))?;
-                Ok(signature)
-            })
+            .ok_or_else(|| anyhow::anyhow!("KMS did not return a signature"))?;
+
+        let decoded_signature = EcdsaSignature::from_der(kms_signature.as_ref())
+            .map_err(|e| anyhow::anyhow!("Failed to decode SubjectPublicKeyInfo: {}", e))?;
+
+        let r_bytes = decoded_signature.r.as_bytes();
+        let s_bytes = decoded_signature.s.as_bytes();
+        let mut signature = [0u8; 65];
+        signature[(32 - r_bytes.len())..32].copy_from_slice(r_bytes);
+        signature[(64 - s_bytes.len())..64].copy_from_slice(decoded_signature.s.as_bytes());
+
+        let public_key = self.get_public_key().await?;
+        for raw_id in 0..4 {
+            let secp = Secp256k1::new();
+            let recid = RecoveryId::try_from(raw_id)
+                .map_err(|e| anyhow::anyhow!("Failed to create RecoveryId: {}", e))?;
+            if let Ok(recovered_public_key) = secp.recover_ecdsa(
+                &Message::from_digest(data),
+                &RecoverableSignature::from_compact(&signature[..64], recid)
+                    .map_err(|e| anyhow::anyhow!("Failed to create RecoverableSignature: {}", e))?,
+            ) {
+                if recovered_public_key == public_key.0 {
+                    signature[64] = raw_id as u8;
+                    return Ok(signature);
+                }
+            }
+        }
+        Err(anyhow::anyhow!(
+            "Failed to recover public key from signature"
+        ))
     }
 
     async fn get_public_key(&self) -> anyhow::Result<(PublicKey, [u8; 20])> {
+        if let Some((public_key, pubkey_evm)) = &self.public_key {
+            return Ok((*public_key, *pubkey_evm));
+        }
+
         let result = self
             .client
             .get_public_key()
@@ -155,9 +228,15 @@ impl Signer for KMSSigner {
         let public_key = result
             .public_key
             .ok_or(anyhow::anyhow!("KMS did not return a public key"))?;
-        let public_key = PublicKey::from_slice(public_key.as_ref())
-            .map_err(|e| anyhow::anyhow!("Failed to create PublicKey from KMS: {}", e))?;
+        let decoded_algorithm_identifier = SubjectPublicKeyInfo::from_der(public_key.as_ref())
+            .map_err(|e| {
+                anyhow::anyhow!("Failed to decode SubjectPublicKeyInfo from KMS: {}", e)
+            })?;
+        let public_key =
+            PublicKey::from_slice(decoded_algorithm_identifier.subject_public_key.raw_bytes())
+                .map_err(|e| anyhow::anyhow!("Failed to create PublicKey from KMS: {}", e))?;
         let pubkey_evm = get_evm_public_key(&public_key)?;
+
         Ok((public_key, pubkey_evm))
     }
 }