Account for P2PK balances

Redefine `Address`'s `AddressHashEnum` to store plain public keys
instead of public key hashes. This ensures we have the necessary
information (the public key itself) to query for P2PK address balance
from electrum while also having access to the public key hash to query
P2PKH & P2WPKH balances.

`AddressHashEnum::AddressHash` variant is kept because sometimes we
don't have access to the public key of that PKH (e.g. when parsing this
address from a script, it's usually presented as PKH).

Add a test case for p2pk address extraction from script & add a p2pk balance test (on testnet)

mm2src/coins/utxo.rs

@@ -1907,6 +1907,12 @@ pub fn output_script(address: &Address, script_type: ScriptType) -> Script {
     match address.addr_format {
         UtxoAddressFormat::Segwit => Builder::build_p2witness(&address.hash),
         _ => match script_type {
+            ScriptType::P2PK => Builder::build_p2pk(
+                address
+                    .hash
+                    .get_pubkey()
+                    .expect("`get_pubkey` is called on a non-PubKey address"),
+            ),
             ScriptType::P2PKH => Builder::build_p2pkh(&address.hash),
             ScriptType::P2SH => Builder::build_p2sh(&address.hash),
             ScriptType::P2WPKH => Builder::build_p2witness(&address.hash),
@@ -1939,12 +1945,12 @@ pub fn address_by_conf_and_pubkey_str(
     let conf_builder = UtxoConfBuilder::new(conf, &params, coin);
     let utxo_conf = try_s!(conf_builder.build());
     let pubkey_bytes = try_s!(hex::decode(pubkey));
-    let hash = dhash160(&pubkey_bytes);
+    let pubkey = try_s!(Public::from_slice(&pubkey_bytes));
 
     let address = Address {
         prefix: utxo_conf.pub_addr_prefix,
         t_addr_prefix: utxo_conf.pub_t_addr_prefix,
-        hash: hash.into(),
+        hash: pubkey.into(),
         checksum_type: utxo_conf.checksum_type,
         hrp: utxo_conf.bech32_hrp,
         addr_format,

mm2src/coins/utxo/qtum.rs

@@ -157,6 +157,7 @@ pub trait QtumBasedCoin: UtxoCommonOps + MarketCoinOps {
         Address {
             prefix: utxo.conf.pub_addr_prefix,
             t_addr_prefix: utxo.conf.pub_t_addr_prefix,
+            // TODO: Should we use AddressHash or ScriptHash (since it's a contract)?
             hash: AddressHashEnum::AddressHash(address.0.into()),
             checksum_type: utxo.conf.checksum_type,
             hrp: utxo.conf.bech32_hrp.clone(),
@@ -169,6 +170,7 @@ pub trait QtumBasedCoin: UtxoCommonOps + MarketCoinOps {
         contract_addr_from_utxo_addr(my_address).mm_err(Qrc20AddressError::from)
     }
 
+    // NOTE: Same as `utxo_addr_from_contract_addr`?
     fn utxo_address_from_contract_addr(&self, address: H160) -> Address {
         let utxo = self.as_ref();
         Address {
@@ -1313,7 +1315,11 @@ pub fn contract_addr_from_str(addr: &str) -> Result<H160, String> { eth::addr_fr
 
 pub fn contract_addr_from_utxo_addr(address: Address) -> MmResult<H160, ScriptHashTypeNotSupported> {
     match address.hash {
+        // DISCUSS: Do we allow pubkeys in here (or pkhs)
+        AddressHashEnum::PubKey(p) => Ok(p.address_hash().take().into()),
+        // TODO: Either AddressHash or ScriptHash should be the correct one (based on `h` being a hash of a script or pubkey).
         AddressHashEnum::AddressHash(h) => Ok(h.take().into()),
+        AddressHashEnum::ScriptHash(h) => Ok(h.take().into()),
         AddressHashEnum::WitnessScriptHash(_) => MmError::err(ScriptHashTypeNotSupported {
             script_hash_type: "Witness".to_owned(),
         }),

mm2src/coins/utxo/rpc_clients.rs

@@ -2320,12 +2320,25 @@ impl UtxoRpcClientOps for ElectrumClient {
     }
 
     fn display_balance(&self, address: Address, decimals: u8) -> RpcRes<BigDecimal> {
-        let hash = electrum_script_hash(&output_script(&address, ScriptType::P2PKH));
-        let hash_str = hex::encode(hash);
-        Box::new(
-            self.scripthash_get_balance(&hash_str)
-                .map(move |electrum_balance| electrum_balance.to_big_decimal(decimals)),
-        )
+        let this = self.clone();
+        let mut hashes = Vec::new();
+        // If the plain pubkey is available, fetch the balance found in P2PK as well (if any).
+        if address.hash.is_pubkey() {
+            hashes.push(hex::encode(electrum_script_hash(&output_script(&address, ScriptType::P2PK))));
+        }
+        hashes.push(hex::encode(electrum_script_hash(&output_script(&address, ScriptType::P2PKH))));
+
+        let fut = async move {
+            Ok(this
+                .scripthash_get_balances(hashes)
+                .compat()
+                .await?
+                .into_iter()
+                .fold(BigDecimal::from(0), |sum, electrum_balance| {
+                    sum + electrum_balance.to_big_decimal(decimals)
+                }))
+        };
+        Box::new(fut.boxed().compat())
     }
 
     fn display_balances(&self, addresses: Vec<Address>, decimals: u8) -> UtxoRpcFut<Vec<(Address, BigDecimal)>> {

mm2src/coins/utxo/slp.rs

@@ -39,7 +39,6 @@ use futures::compat::Future01CompatExt;
 use futures::{FutureExt, TryFutureExt};
 use futures01::Future;
 use hex::FromHexError;
-use keys::hash::H160;
 use keys::{AddressHashEnum, CashAddrType, CashAddress, CompactSignature, KeyPair, NetworkPrefix as CashAddrPrefix,
            Public};
 use mm2_core::mm_ctx::MmArc;
@@ -388,7 +387,7 @@ impl SlpToken {
             let my_public_key = self.platform_coin.my_public_key()?;
             let slp_change_out = TransactionOutput {
                 value: self.platform_dust(),
-                script_pubkey: ScriptBuilder::build_p2pkh(&my_public_key.address_hash().into()).to_bytes(),
+                script_pubkey: ScriptBuilder::build_p2pkh(&my_public_key.clone().into()).to_bytes(),
             };
             outputs.push(slp_change_out);
         }
@@ -424,7 +423,8 @@ impl SlpToken {
         amount: u64,
     ) -> Result<UtxoTx, TransactionErr> {
         let payment_script = payment_script(time_lock, secret_hash, my_pub, other_pub);
-        let script_pubkey = ScriptBuilder::build_p2sh(&dhash160(&payment_script).into()).to_bytes();
+        let script_hash = AddressHashEnum::ScriptHash(dhash160(&payment_script));
+        let script_pubkey = ScriptBuilder::build_p2sh(&script_hash).to_bytes();
         let slp_out = SlpOutput { amount, script_pubkey };
         let (preimage, recently_spent) = try_tx_s!(self.generate_slp_tx_preimage(vec![slp_out]).await);
         generate_and_send_tx(
@@ -633,7 +633,7 @@ impl SlpToken {
         outputs.push(op_return_out_mm);
 
         let my_public_key = self.platform_coin.my_public_key()?;
-        let my_script_pubkey = ScriptBuilder::build_p2pkh(&my_public_key.address_hash().into());
+        let my_script_pubkey = ScriptBuilder::build_p2pkh(&my_public_key.clone().into());
         let slp_output = TransactionOutput {
             value: self.platform_dust(),
             script_pubkey: my_script_pubkey.to_bytes(),
@@ -1203,7 +1203,7 @@ impl SwapOps for SlpToken {
     fn send_taker_fee(&self, fee_addr: &[u8], dex_fee: DexFee, _uuid: &[u8]) -> TransactionFut {
         let coin = self.clone();
         let fee_pubkey = try_tx_fus!(Public::from_slice(fee_addr));
-        let script_pubkey = ScriptBuilder::build_p2pkh(&fee_pubkey.address_hash().into()).into();
+        let script_pubkey = ScriptBuilder::build_p2pkh(&fee_pubkey.into()).into();
         let amount = try_tx_fus!(dex_fee.fee_uamount(self.decimals()));
 
         let fut = async move {
@@ -1784,7 +1784,7 @@ impl MmCoin for SlpToken {
             },
         };
         // can use dummy P2SH script_pubkey here
-        let script_pubkey = ScriptBuilder::build_p2sh(&H160::default().into()).into();
+        let script_pubkey = ScriptBuilder::build_p2sh(&AddressHashEnum::default_script_hash()).into();
         let slp_out = SlpOutput {
             amount: slp_amount,
             script_pubkey,
@@ -1833,7 +1833,7 @@ impl MmCoin for SlpToken {
     ) -> TradePreimageResult<TradeFee> {
         let slp_amount = sat_from_big_decimal(&dex_fee_amount.fee_amount().into(), self.decimals())?;
         // can use dummy P2PKH script_pubkey here
-        let script_pubkey = ScriptBuilder::build_p2pkh(&H160::default().into()).into();
+        let script_pubkey = ScriptBuilder::build_p2pkh(&AddressHashEnum::default_address_hash()).into();
         let slp_out = SlpOutput {
             amount: slp_amount,
             script_pubkey,
@@ -2158,7 +2158,7 @@ mod slp_tests {
 
         let invalid_slp_send_out = TransactionOutput {
             value: 1000,
-            script_pubkey: ScriptBuilder::build_p2sh(&dhash160(&htlc_script).into()).into(),
+            script_pubkey: ScriptBuilder::build_p2sh(&AddressHashEnum::ScriptHash(dhash160(&htlc_script))).into(),
         };
 
         let tx_err = block_on(generate_and_send_tx(

mm2src/coins/utxo/utxo_builder/utxo_coin_builder.rs

@@ -232,7 +232,7 @@ where
     let my_address = Address {
         prefix: conf.pub_addr_prefix,
         t_addr_prefix: conf.pub_t_addr_prefix,
-        hash: AddressHashEnum::AddressHash(key_pair.public().address_hash()),
+        hash: key_pair.public().clone().into(),
         checksum_type: conf.checksum_type,
         hrp: conf.bech32_hrp.clone(),
         addr_format,

mm2src/coins/utxo/utxo_common.rs

@@ -629,7 +629,7 @@ pub fn addresses_from_script<T: UtxoCommonOps>(coin: &T, script: &Script) -> Res
         .into_iter()
         .map(|dst| {
             let (prefix, t_addr_prefix, addr_format) = match dst.kind {
-                ScriptType::P2PKH => (
+                ScriptType::P2PKH | ScriptType::P2PK => (
                     conf.pub_addr_prefix,
                     conf.pub_t_addr_prefix,
                     coin.addr_format_for_standard_scripts(),
@@ -1539,7 +1539,7 @@ pub async fn sign_and_send_taker_funding_spend<T: UtxoCommonOps>(
         );
         let payment_address = Address {
             checksum_type: coin.as_ref().conf.checksum_type,
-            hash: AddressHashEnum::AddressHash(dhash160(&payment_redeem_script)),
+            hash: AddressHashEnum::ScriptHash(dhash160(&payment_redeem_script)),
             prefix: coin.as_ref().conf.p2sh_addr_prefix,
             t_addr_prefix: coin.as_ref().conf.p2sh_t_addr_prefix,
             hrp: coin.as_ref().conf.bech32_hrp.clone(),
@@ -2606,7 +2606,8 @@ pub fn watcher_validate_taker_payment<T: UtxoCommonOps + SwapOps>(
             },
         };
 
-        if taker_payment_locking_script != Builder::build_p2sh(&dhash160(&expected_redeem).into()).to_bytes() {
+        let script_hash = AddressHashEnum::ScriptHash(dhash160(&expected_redeem));
+        if taker_payment_locking_script != Builder::build_p2sh(&script_hash).to_bytes() {
             return MmError::err(ValidatePaymentError::WrongPaymentTx(format!(
                 "{INVALID_SCRIPT_ERR_LOG}: Payment tx locking script {taker_payment_locking_script:?} doesn't match expected"
             )));
@@ -2718,8 +2719,8 @@ pub fn check_if_my_payment_sent<T: UtxoCommonOps + SwapOps>(
         my_htlc_keypair.public(),
         &try_fus!(Public::from_slice(other_pub)),
     );
-    let hash = dhash160(&script);
-    let p2sh = Builder::build_p2sh(&hash.into());
+    let hash = AddressHashEnum::ScriptHash(dhash160(&script));
+    let p2sh = Builder::build_p2sh(&hash);
     let script_hash = electrum_script_hash(&p2sh);
     let fut = async move {
         match &coin.as_ref().rpc_client {
@@ -2739,7 +2740,7 @@ pub fn check_if_my_payment_sent<T: UtxoCommonOps + SwapOps>(
                 let target_addr = Address {
                     t_addr_prefix: coin.as_ref().conf.p2sh_t_addr_prefix,
                     prefix: coin.as_ref().conf.p2sh_addr_prefix,
-                    hash: hash.into(),
+                    hash,
                     checksum_type: coin.as_ref().conf.checksum_type,
                     hrp: coin.as_ref().conf.bech32_hrp.clone(),
                     addr_format: coin.addr_format().clone(),
@@ -2937,7 +2938,7 @@ pub fn verify_message<T: UtxoCommonOps>(
     let signature = CompactSignature::from(base64::decode(signature_base64)?);
     let recovered_pubkey = Public::recover_compact(&H256::from(message_hash), &signature)?;
     let received_address = checked_address_from_str(coin, address)?;
-    Ok(AddressHashEnum::from(recovered_pubkey.address_hash()) == received_address.hash)
+    Ok(received_address.hash == recovered_pubkey.into())
 }
 
 pub fn my_balance<T>(coin: T) -> BalanceFut<CoinBalance>
@@ -4586,7 +4587,7 @@ pub fn address_from_raw_pubkey(
     Ok(Address {
         t_addr_prefix,
         prefix,
-        hash: try_s!(Public::from_slice(pub_key)).address_hash().into(),
+        hash: try_s!(Public::from_slice(pub_key)).into(),
         checksum_type,
         hrp,
         addr_format,
@@ -4604,7 +4605,7 @@ pub fn address_from_pubkey(
     Address {
         t_addr_prefix,
         prefix,
-        hash: pub_key.address_hash().into(),
+        hash: pub_key.clone().into(),
         checksum_type,
         hrp,
         addr_format,
@@ -4652,7 +4653,8 @@ pub fn validate_payment<T: UtxoCommonOps>(
             )));
         }
 
-        let expected_script_pubkey: Bytes = Builder::build_p2sh(&dhash160(&expected_redeem).into()).into();
+        let script_hash = AddressHashEnum::ScriptHash(dhash160(&expected_redeem));
+        let expected_script_pubkey: Bytes = Builder::build_p2sh(&script_hash).into();
 
         let actual_output = match tx.outputs.get(output_index) {
             Some(output) => output,
@@ -4710,7 +4712,8 @@ async fn search_for_swap_output_spend(
         return ERR!("Transaction doesn't have any output");
     }
     let script = payment_script(time_lock, secret_hash, first_pub, second_pub);
-    let expected_script_pubkey = Builder::build_p2sh(&dhash160(&script).into()).to_bytes();
+    let script_hash = AddressHashEnum::ScriptHash(dhash160(&script));
+    let expected_script_pubkey = Builder::build_p2sh(&script_hash).to_bytes();
     let script_pubkey = &tx
         .outputs
         .get(output_index)
@@ -4797,11 +4800,11 @@ where
             swap_proto_v2_scripts::taker_payment_script(time_lock, secret_hash, &my_public, &other_public)
         },
     };
-    let redeem_script_hash = dhash160(&redeem_script);
+    let redeem_script_hash = AddressHashEnum::ScriptHash(dhash160(&redeem_script));
     let amount = try_s!(sat_from_big_decimal(&amount, coin.as_ref().decimals));
     let htlc_out = TransactionOutput {
         value: amount,
-        script_pubkey: Builder::build_p2sh(&redeem_script_hash.into()).into(),
+        script_pubkey: Builder::build_p2sh(&redeem_script_hash).into(),
     };
     // record secret hash to blockchain too making it impossible to lose
     // lock time may be easily brute forced so it is not mandatory to record it
@@ -4823,7 +4826,7 @@ where
 
     let payment_address = Address {
         checksum_type: coin.as_ref().conf.checksum_type,
-        hash: redeem_script_hash.into(),
+        hash: redeem_script_hash,
         prefix: coin.as_ref().conf.p2sh_addr_prefix,
         t_addr_prefix: coin.as_ref().conf.p2sh_t_addr_prefix,
         hrp: coin.as_ref().conf.bech32_hrp.clone(),

mm2src/coins/utxo/utxo_common_tests.rs

@@ -63,7 +63,7 @@ pub(super) fn utxo_coin_fields_for_test(
     let key_pair = key_pair_from_seed(&seed).unwrap();
     let my_address = Address {
         prefix: 60,
-        hash: key_pair.public().address_hash().into(),
+        hash: key_pair.public().clone().into(),
         t_addr_prefix: 0,
         checksum_type,
         hrp: if is_segwit_coin {

mm2src/coins/utxo/utxo_tests.rs

@@ -281,17 +281,30 @@ fn test_generate_transaction() {
 fn test_addresses_from_script() {
     let client = electrum_client_for_test(DOC_ELECTRUM_ADDRS);
     let coin = utxo_coin_for_test(client.into(), None, false);
+    // P2PK
+    let script: Script = "2102b4c3132caa3c998cf98abeaf05226f93565621c6f11038688743c6aa7bb142b3ac".into();
+    let expected_addr = Address {
+        prefix: coin.as_ref().conf.pub_addr_prefix,
+        t_addr_prefix: coin.as_ref().conf.pub_t_addr_prefix,
+        hrp: None,
+        hash: Public::Compressed("02b4c3132caa3c998cf98abeaf05226f93565621c6f11038688743c6aa7bb142b3".into()).into(),
+        checksum_type: coin.as_ref().conf.checksum_type,
+        addr_format: coin.addr_format_for_standard_scripts(),
+    };
+    let actual_addr = coin.addresses_from_script(&script).unwrap();
+    assert_eq!(vec![expected_addr], actual_addr);
+
     // P2PKH
     let script: Script = "76a91405aab5342166f8594baf17a7d9bef5d56744332788ac".into();
-    let expected_addr: Vec<Address> = vec!["R9o9xTocqr6CeEDGDH6mEYpwLoMz6jNjMW".into()];
+    let expected_addr = Address::p2pkh_address("R9o9xTocqr6CeEDGDH6mEYpwLoMz6jNjMW").unwrap();
     let actual_addr = coin.addresses_from_script(&script).unwrap();
-    assert_eq!(expected_addr, actual_addr);
+    assert_eq!(vec![expected_addr], actual_addr);
 
     // P2SH
     let script: Script = "a914e71a6120653ebd526e0f9d7a29cde5969db362d487".into();
-    let expected_addr: Vec<Address> = vec!["bZoEPR7DjTqSDiQTeRFNDJuQPTRY2335LD".into()];
+    let expected_addr = Address::p2sh_address("bZoEPR7DjTqSDiQTeRFNDJuQPTRY2335LD").unwrap();
     let actual_addr = coin.addresses_from_script(&script).unwrap();
-    assert_eq!(expected_addr, actual_addr);
+    assert_eq!(vec![expected_addr], actual_addr);
 }
 
 #[test]

mm2src/coins/utxo/utxo_withdraw.rs

@@ -12,7 +12,7 @@ use crypto::hw_rpc_task::HwRpcTaskAwaitingStatus;
 use crypto::trezor::trezor_rpc_task::{TrezorRequestStatuses, TrezorRpcTaskProcessor};
 use crypto::trezor::{TrezorError, TrezorProcessingError};
 use crypto::{from_hw_error, CryptoCtx, CryptoCtxError, DerivationPath, HwError, HwProcessingError, HwRpcError};
-use keys::{AddressFormat, AddressHashEnum, KeyPair, Private, Public as PublicKey, Type as ScriptType};
+use keys::{AddressFormat, KeyPair, Private, Public as PublicKey, Type as ScriptType};
 use mm2_core::mm_ctx::MmArc;
 use mm2_err_handle::prelude::*;
 use rpc::v1::types::ToTxHash;
@@ -468,7 +468,7 @@ where
                 let my_address = Address {
                     prefix: coin.as_ref().conf.pub_addr_prefix,
                     t_addr_prefix: coin.as_ref().conf.pub_t_addr_prefix,
-                    hash: AddressHashEnum::AddressHash(key_pair.public().address_hash()),
+                    hash: key_pair.public().clone().into(),
                     checksum_type: coin.as_ref().conf.checksum_type,
                     hrp: coin.as_ref().conf.bech32_hrp.clone(),
                     addr_format,

mm2src/coins/utxo_signer/src/with_key_pair.rs

@@ -109,7 +109,7 @@ pub fn p2pkh_spend(
 ) -> UtxoSignWithKeyPairResult<TransactionInput> {
     let unsigned_input = get_input(signer, input_index)?;
 
-    let script = Builder::build_p2pkh(&key_pair.public().address_hash().into());
+    let script = Builder::build_p2pkh(&key_pair.public().clone().into());
     if script != prev_script {
         return MmError::err(UtxoSignWithKeyPairError::MismatchScript {
             script_type: "P2PKH".to_owned(),
@@ -176,8 +176,8 @@ pub fn p2wpkh_spend(
 ) -> UtxoSignWithKeyPairResult<TransactionInput> {
     let unsigned_input = get_input(signer, input_index)?;
 
-    let script_code = Builder::build_p2pkh(&key_pair.public().address_hash().into()); // this is the scriptCode by BIP-0143: for P2WPKH scriptCode is P2PKH
-    let script_pub_key = Builder::build_p2witness(&key_pair.public().address_hash().into());
+    let script_code = Builder::build_p2pkh(&key_pair.public().clone().into()); // this is the scriptCode by BIP-0143: for P2WPKH scriptCode is P2PKH
+    let script_pub_key = Builder::build_p2witness(&key_pair.public().clone().into());
     if script_pub_key != prev_script {
         return MmError::err(UtxoSignWithKeyPairError::MismatchScript {
             script_type: "P2WPKH".to_owned(),

mm2src/coins/z_coin/z_htlc.rs

@@ -16,7 +16,7 @@ use crate::{PrivKeyPolicyNotAllowed, TransactionEnum};
 use bitcrypto::dhash160;
 use derive_more::Display;
 use futures::compat::Future01CompatExt;
-use keys::Address;
+use keys::{Address, AddressHashEnum};
 use keys::{KeyPair, Public};
 use mm2_err_handle::prelude::*;
 use mm2_number::BigDecimal;
@@ -46,11 +46,11 @@ pub async fn z_send_htlc(
     amount: BigDecimal,
 ) -> Result<ZTransaction, MmError<SendOutputsErr>> {
     let payment_script = payment_script(time_lock, secret_hash, my_pub, other_pub);
-    let script_hash = dhash160(&payment_script);
+    let script_hash = AddressHashEnum::ScriptHash(dhash160(&payment_script));
     let htlc_address = Address {
         prefix: coin.utxo_arc.conf.p2sh_addr_prefix,
         t_addr_prefix: coin.utxo_arc.conf.p2sh_t_addr_prefix,
-        hash: script_hash.into(),
+        hash: script_hash.clone(),
         checksum_type: coin.utxo_arc.conf.checksum_type,
         addr_format: UtxoAddressFormat::Standard,
         hrp: None,