Merge branch 'main' into hashicorp-feature-tmkms

src/privval.rs

@@ -129,6 +129,38 @@ impl SignableMsg {
             },
         }
     }
+
+    /// Add a consensus signature to this message.
+    pub fn add_consensus_signature(&mut self, signature: impl Into<tendermint::Signature>) {
+        match self {
+            SignableMsg::Proposal(proposal) => {
+                proposal.signature = Some(signature.into());
+            }
+            SignableMsg::Vote(vote) => {
+                vote.signature = Some(signature.into());
+            }
+        }
+    }
+
+    /// Add an extension signature to this message.
+    pub fn add_extension_signature(
+        &mut self,
+        signature: impl Into<tendermint::Signature>,
+    ) -> Result<(), Error> {
+        match self {
+            SignableMsg::Vote(vote) => {
+                if vote.extension.is_empty() {
+                    return Err(Error::protocol(
+                        "can't add signature to empty extension".into(),
+                    ));
+                }
+
+                vote.extension_signature = Some(signature.into());
+                Ok(())
+            }
+            _ => Err(Error::invalid_message_type()),
+        }
+    }
 }
 
 impl From<Proposal> for SignableMsg {

src/rpc.rs

@@ -3,7 +3,7 @@
 // TODO: docs for everything
 #![allow(missing_docs)]
 
-use crate::{keyring::Signature, privval::SignableMsg};
+use crate::privval::SignableMsg;
 use prost::Message as _;
 use std::io::Read;
 use tendermint::{chain, Proposal, Vote};
@@ -128,28 +128,21 @@ impl Response {
             }),
         }
     }
+}
 
-    /// Construct a signed response from a [`SignableMsg`] and a [`Signature`].
-    pub fn sign(msg: SignableMsg, sig: Signature) -> Result<Response, Error> {
+impl From<SignableMsg> for Response {
+    fn from(msg: SignableMsg) -> Response {
         match msg {
             SignableMsg::Proposal(proposal) => {
-                let mut proposal = proto::types::Proposal::from(proposal);
-                proposal.signature = sig.to_vec();
-                Ok(Response::SignedProposal(
-                    proto::privval::SignedProposalResponse {
-                        proposal: Some(proposal),
-                        error: None,
-                    },
-                ))
-            }
-            SignableMsg::Vote(vote) => {
-                let mut vote = proto::types::Vote::from(vote);
-                vote.signature = sig.to_vec();
-                Ok(Response::SignedVote(proto::privval::SignedVoteResponse {
-                    vote: Some(vote),
+                Response::SignedProposal(proto::privval::SignedProposalResponse {
+                    proposal: Some(proposal.into()),
                     error: None,
-                }))
+                })
             }
+            SignableMsg::Vote(vote) => Response::SignedVote(proto::privval::SignedVoteResponse {
+                vote: Some(vote.into()),
+                error: None,
+            }),
         }
     }
 }

src/session.rs

@@ -142,23 +142,28 @@ impl Session {
         // TODO(tarcieri): support for non-default public keys
         let public_key = None;
         let chain_id = self.config.chain_id.clone();
-
         let canonical_msg = signable_msg.canonical_bytes(chain_id.clone())?;
+
         let started_at = Instant::now();
-        let signature = chain.keyring.sign(public_key, &canonical_msg)?;
+        let consensus_sig = chain.keyring.sign(public_key, &canonical_msg)?;
+        signable_msg.add_consensus_signature(consensus_sig);
         self.log_signing_request(&signable_msg, started_at).unwrap();
 
         // Add extension signature if there are any extensions defined
         if let Some(extension_msg) = signable_msg.extension_bytes(chain_id)? {
+            let started_at = Instant::now();
             let extension_sig = chain.keyring.sign(public_key, &extension_msg)?;
-
-            match &mut signable_msg {
-                SignableMsg::Vote(vote) => vote.extension_signature = Some(extension_sig.into()),
-                other => fail!(InvalidMessageError, "expected a vote type: {:?}", other),
-            }
+            signable_msg.add_extension_signature(extension_sig)?;
+
+            info!(
+                "[{}@{}] signed vote extension ({} ms)",
+                &self.config.chain_id,
+                &self.config.addr,
+                started_at.elapsed().as_millis(),
+            );
         }
 
-        Response::sign(signable_msg, signature)
+        Ok(signable_msg.into())
     }
 
     /// If a max block height is configured, ensure the block we're signing