kms_secp256k1_api/services/mocks/
mock_cosmos_keys_service.rs

1use crate::{
2    config::Config,
3    constants::{COSMOS_SECP_LEN, DEFAULT_COSMOS_HRP},
4    services::{
5        cosmos_keys_service::{
6            BodyHelper, append_signature_to_transaction, build_auth_info, fee_amount_json,
7            fetch_account_info, signature_to_json,
8        },
9        crypto_service::CryptoService,
10        keys_service::{KeyEntry, KeysServiceTrait, SigEntry},
11        mocks::mock_keys_service::{KeyPair, MockKeysService},
12    },
13};
14use base64::Engine;
15use base64::engine::general_purpose::STANDARD;
16use cosmrs::{
17    bip32::{PrivateKey, PublicKey},
18    proto::cosmos::tx::v1beta1::TxRaw,
19    tendermint::chain::Id,
20    tx::{Fee, MessageExt, SignDoc},
21};
22use k256::{
23    ecdsa::{Signature, SigningKey, signature::Signer},
24    elliptic_curve::rand_core::OsRng,
25    sha2::{Digest, Sha256},
26};
27use serde_json::json;
28use std::str::FromStr;
29use tracing::error;
30
31pub struct MockCosmosKeysService {
32    inner: MockKeysService,
33    hrp: String,
34}
35
36#[async_trait::async_trait]
37impl KeysServiceTrait for MockCosmosKeysService {
38    async fn create_key(&mut self, _config: &Config) -> Result<KeyEntry, String> {
39        let signing_key = SigningKey::random(&mut OsRng);
40        let verifying_key = signing_key.verifying_key();
41        let encoded_point = verifying_key.to_encoded_point(true);
42
43        let private_key_bytes = signing_key.to_bytes();
44        let private_key = hex::encode(private_key_bytes);
45        let pubkey_bytes = encoded_point.as_bytes();
46        let public_key = hex::encode(pubkey_bytes);
47        let address = self.resolve_key(&public_key)?;
48
49        {
50            let key_pair = KeyPair {
51                public_key: public_key.clone(),
52                private_key,
53                address: address.clone(),
54            };
55            let mut keys = self.inner.keys.lock().await;
56            keys.entry(address.clone()).or_insert(key_pair);
57        }
58
59        Ok(KeyEntry {
60            public_key: Some(public_key.clone()).into(),
61            address: address.clone().into(),
62            public_key_base64: STANDARD.encode(public_key).into(),
63            key_id: address.into(),
64        })
65    }
66
67    /// Signs a given transaction hash using the specified key.
68    ///
69    /// Takes the transaction hash, key, and configuration.
70    /// Returns the signature as a hexadecimal string if successful.
71    /// Returns an error string if signing fails.
72    async fn sign_transaction_hash(
73        &mut self,
74        _config: &Config,
75        transaction_hash: &str,
76        key: &str,
77    ) -> Result<SigEntry, String> {
78        let key = self.resolve_key(key)?;
79
80        let key_pair = {
81            let keys = self.inner.keys.lock().await;
82            keys.get(&key)
83                .ok_or_else(|| "Key not found".to_string())?
84                .clone()
85        };
86
87        let hash_bytes = hex::decode(transaction_hash)
88            .map_err(|e| format!("Invalid transaction hash hex: {e}"))?;
89
90        if hash_bytes.len() != 32 {
91            return Err("Transaction hash must be 32 bytes".to_string());
92        }
93
94        let private_key_bytes = hex::decode(&key_pair.private_key)
95            .map_err(|e| format!("Failed to decode secret key: {e}"))?;
96
97        let signing_key = SigningKey::from_slice(&private_key_bytes)
98            .map_err(|e| format!("Failed to create signing key: {e}"))?;
99
100        let signature: Signature = signing_key.sign(&hash_bytes);
101        let der = signature.to_der();
102        let der_bytes = der.as_bytes();
103        let base64_signature = STANDARD.encode(der_bytes);
104        // let sig_hex_long = hex::encode(der_bytes);
105
106        let signature = self
107            .inner
108            .crypto_service
109            .convert(&base64_signature)
110            .map_err(|e| {
111                let msg = format!("Failed to convert signature: {e}");
112                error!("{}", msg);
113                msg
114            })?;
115
116        // Verify signature
117        let is_valid = self
118            .verify(transaction_hash, &signature, &key_pair.public_key)
119            .await?;
120
121        if !is_valid {
122            return Err("Signature verification failed".to_string());
123        }
124
125        Ok(SigEntry {
126            address: key.into(),
127            public_key: key_pair.public_key.into(),
128            signature: signature.into(),
129        })
130    }
131
132    /// Signs a transaction represented as a JSON string with the given public key.
133    ///
134    /// Returns the signed transaction as a JSON string on success.
135    /// Returns an error string if parsing the transaction or signing fails.
136    #[allow(clippy::too_many_lines)]
137    async fn sign_transaction(
138        &mut self,
139        config: &Config,
140        transaction_str: &str,
141        key: &str,
142    ) -> Result<String, String> {
143        let tx_json: serde_json::Value = serde_json::from_str(transaction_str)
144            .map_err(|e| format!("Failed to parse JSON: {e}"))?;
145
146        // Extract chain_id and address
147        let chain_id = Id::from_str(&config.get_cosmos_chain_id())
148            .map_err(|e| format!("Failed to fetch chain_id: {e}"))?;
149
150        // Get key pair and wallet
151        let key = self.resolve_key(key)?;
152        let key_pair = {
153            let keys = self.inner.keys.lock().await;
154            keys.get(&key)
155                .ok_or_else(|| "Public key not found".to_string())?
156                .clone()
157        };
158        let key = key_pair.address;
159
160        // Deserialize TxBody
161        let helper: BodyHelper = serde_json::from_value(tx_json["body"].clone())
162            .map_err(|e| format!("Invalid TxBody: {e}"))?;
163
164        let tx_body = helper.into_body()?;
165
166        // Deserialize Fee
167        let fee: Fee = serde_json::from_value(tx_json["auth_info"]["fee"].clone())
168            .map_err(|e| format!("Invalid Fee: {e}"))?;
169
170        let private_key_bytes = hex::decode(&key_pair.private_key)
171            .map_err(|e| format!("Failed to decode secret key: {e}"))?;
172
173        let signing_key = SigningKey::from_slice(&private_key_bytes)
174            .map_err(|e| format!("Failed to create signing key: {e}"))?;
175
176        let public_key = signing_key.public_key();
177
178        let public_key_bytes = public_key.to_bytes();
179        let pubkey_base64 = STANDARD.encode(public_key_bytes);
180
181        // Fetch account_number and sequence
182        let mut account = fetch_account_info(&key, &pubkey_base64, config)
183            .await
184            .map_err(|e| format!("Failed to fetch account info: {e}"))?;
185
186        let Some(fetched_pub_key) = account.pub_key.take() else {
187            return Err("Account info is missing pub_key".to_string());
188        };
189
190        if fetched_pub_key.key != pubkey_base64 {
191            return Err(format!(
192                "Invalid fetched public key: got {}",
193                fetched_pub_key.key
194            ));
195        }
196        let auth_info = build_auth_info(&public_key_bytes, account.sequence, &fee)?;
197
198        // Build SignDoc
199        let sign_doc = SignDoc::new(&tx_body, &auth_info, &chain_id, account.sequence)
200            .map_err(|e| format!("SignDoc error: {e}"))?;
201
202        let sign_doc_bytes = sign_doc
203            .into_bytes()
204            .map_err(|e| format!("SignDoc encode error: {e}"))?;
205
206        // Hash and sign
207        let transaction_hash = Sha256::digest(&sign_doc_bytes);
208        let signature: Signature = signing_key.sign(&transaction_hash);
209        let transaction_hash = hex::encode(transaction_hash);
210
211        // Verify signature
212        let signature_hex = signature.to_string();
213        let public_key = hex::encode(public_key_bytes);
214
215        let is_valid = self
216            .verify(&transaction_hash, &signature_hex, &public_key)
217            .await?;
218        if !is_valid {
219            return Err("Generated signature failed verification".to_string());
220        }
221
222        let body_bytes = tx_body
223            .into_bytes()
224            .map_err(|e| format!("Failed to encode body: {e}"))?;
225        let auth_info_bytes = auth_info
226            .into_bytes()
227            .map_err(|e| format!("Failed to encode auth_info: {e}"))?;
228
229        let tx_raw = TxRaw {
230            body_bytes,
231            auth_info_bytes,
232            signatures: vec![signature.to_bytes().to_vec()],
233        };
234
235        let tx_raw_bytes = tx_raw
236            .to_bytes()
237            .map_err(|e| format!("Failed to encode TxRaw: {e}"))?;
238
239        let base64_tx = STANDARD.encode(tx_raw_bytes);
240        let broadcast_request = json!({
241            "tx_bytes": base64_tx,
242            "mode": "BROADCAST_MODE_SYNC"  // or "BLOCK" or "ASYNC"
243        });
244        let fee_amount_json = fee_amount_json(&fee);
245
246        // Existing signatures from the original transaction (if any)
247        let new_signature = signature_to_json(&key, &public_key, &signature, &transaction_hash);
248        let signatures_array = append_signature_to_transaction(&tx_json, new_signature);
249
250        // Prepare final JSON response
251        let result = json!({
252            "chain_id": chain_id,
253            "body": tx_json["body"],
254            "auth_info": {
255                "signer_infos": [
256                    {
257                        "public_key": {
258                            "@type": fetched_pub_key.key_type,
259                            "key": fetched_pub_key.key,
260                        },
261                        "mode_info": {
262                            "single": { "mode": "SIGN_MODE_DIRECT" }
263                        },
264                        "sequence": account.sequence,
265                        "account_number": account.account_number
266                    }
267                ],
268                "fee": {
269                    "amount": fee_amount_json,
270                    "gas_limit": fee.gas_limit
271                }
272            },
273            "broadcast_request": broadcast_request,
274            "signatures": signatures_array
275        });
276
277        // Return the wrapped transaction + signatures JSON
278        serde_json::to_string(&result)
279            .map_err(|e| format!("Failed to serialize final signed transaction: {e}"))
280    }
281
282    /// Verifies an Cosmos EIP-155 signature for a given transaction hash and public key.
283    ///
284    /// Returns `Ok(true)` if the signature is valid, `Ok(false)` if invalid.
285    /// Returns an error string for failures such as invalid formats.
286    async fn verify(
287        &mut self,
288        transaction_hash_hex: &str,
289        signature_hex: &str,
290        key: &str,
291    ) -> Result<bool, String> {
292        let key = self.resolve_key(key)?;
293        let key_pair = {
294            let keys = self.inner.keys.lock().await;
295            keys.get(&key)
296                .ok_or_else(|| "Public key not found".to_string())?
297                .clone()
298        };
299        self.inner
300            .verify(transaction_hash_hex, signature_hex, &key_pair.public_key)
301            .map_err(|e| {
302                let msg = format!("Signature verification failed: {e}");
303                error!("{}", msg);
304                msg
305            })
306    }
307
308    /// Verifies a signature via the Key Management Service (KMS).
309    ///
310    /// Asynchronous function that returns `Ok(true)` if verification succeeds,
311    /// or an error string if it fails.
312    async fn verify_via_kms(
313        &mut self,
314        transaction_hash_hex: &str,
315        signature_hex: &str,
316        key: &str,
317    ) -> Result<bool, String> {
318        let key = self.resolve_key(key)?;
319        let key_pair = {
320            let keys = self.inner.keys.lock().await;
321            keys.get(&key)
322                .ok_or_else(|| "Public key not found".to_string())?
323                .clone()
324        };
325        self.inner
326            .verify_via_kms(transaction_hash_hex, signature_hex, &key_pair.public_key)
327            .await
328    }
329
330    /// Deletes a public key from storage.
331    ///
332    /// Returns `Ok(true)` if the key was deleted, `Ok(false)` if the key was not found.
333    /// Returns an error string if deletion fails.
334    async fn delete_key(&mut self, key: &str) -> Result<bool, String> {
335        let key = self.resolve_key(key)?;
336        Ok(self.inner.delete_key(&key).await)
337    }
338
339    /// Lists all stored keys along with their associated metadata.
340    ///
341    /// Returns a vector of `KeyEntry` on success.
342    /// Returns an error string on failure.
343    async fn list_keys(&mut self) -> Result<Vec<KeyEntry>, String> {
344        Ok(self.inner.list_keys().await)
345    }
346}
347
348impl MockCosmosKeysService {
349    /// Creates a new `MockCosmosKeysService` with an empty in-memory key store.
350    ///
351    /// # Arguments
352    ///
353    /// * `_config` - Unused configuration object, included for interface compatibility.
354    /// * `crypto_service` - The crypto service to use for key operations.
355    ///
356    /// # Errors
357    ///
358    /// This function currently does not return an error, but it returns a `Result`
359    /// to match a common interface and allow future fallibility.
360    pub fn new(config: Config, crypto_service: CryptoService) -> Result<Self, String> {
361        let cosmos_hrp = config.get_cosmos_hrp();
362        let hrp = match cosmos_hrp.as_str() {
363            "" => DEFAULT_COSMOS_HRP,
364            hrp => hrp,
365        }
366        .to_string();
367
368        Ok(Self {
369            inner: MockKeysService::new(config, crypto_service),
370            hrp,
371        })
372    }
373
374    /// Resolves a given key string to a Cosmos Bech32 address.
375    ///
376    /// This function checks whether the input `key` is a compressed secp256k1 public key
377    /// (by comparing its length to the expected `COSMOS_SECP_LEN`). If so, it attempts to convert
378    /// the public key to its corresponding Cosmos address using the underlying crypto service and the configured `hrp`.
379    /// Otherwise, it assumes the key is already a valid address and returns it unchanged.
380    ///
381    /// # Parameters
382    /// - `key`: A string that is either a Cosmos Bech32 address or a compressed public key (hex-encoded, 33 bytes).
383    ///
384    /// # Returns
385    /// - `Ok(String)`: The resolved Cosmos address as a Bech32-encoded string.
386    /// - `Err(String)`: An error message if the conversion from public key to address fails.
387    ///
388    /// # Errors
389    /// - Returns an error if the input is assumed to be a public key and the address derivation fails.
390    ///
391    /// # Notes
392    /// - The Bech32 address is generated using the provided `hrp` as prefix.
393    /// - This function logs an error internally if conversion fails.
394    fn resolve_key(&mut self, key: &str) -> Result<String, String> {
395        if key.len() == COSMOS_SECP_LEN {
396            self.inner
397                .crypto_service
398                .address_cosmos(key, &self.hrp)
399                .map_err(|e| {
400                    let msg = format!("Failed to convert public key to address: {e:?}");
401                    error!("{}", &msg);
402                    msg
403                })
404        } else {
405            Ok(key.to_string())
406        }
407    }
408}
kms_secp256k1_api/services/mocks/mock_cosmos_keys_service.rs

kms_secp256k1_api/services/mocks/
mock_cosmos_keys_service.rs
