account for p2pk balances

This adds a new `.pubkey` field `Address` of type `Option<Public>`.
Using this, we have a way to keep the pubkey of our address so that we
can pull p2pk balance if any.

mm2src/coins/qrc20/qrc20_tests.rs

@@ -68,12 +68,11 @@ fn test_withdraw_to_p2sh_address_should_fail() {
 
     let p2sh_address = AddressBuilder::new(
         UtxoAddressFormat::Standard,
-        coin.as_ref().derivation_method.unwrap_single_addr().hash().clone(),
         *coin.as_ref().derivation_method.unwrap_single_addr().checksum_type(),
         coin.as_ref().conf.address_prefixes.clone(),
         coin.as_ref().conf.bech32_hrp.clone(),
     )
-    .as_sh()
+    .as_sh(coin.as_ref().derivation_method.unwrap_single_addr().hash().clone())
     .build()
     .expect("valid address props");
 

mm2src/coins/utxo.rs

@@ -1925,6 +1925,8 @@ pub fn output_script(address: &Address) -> Result<Script, keys::Error> {
     }
 }
 
+pub fn output_script_p2pk(pubkey: &Public) -> Script { Builder::build_p2pk(pubkey) }
+
 pub fn address_by_conf_and_pubkey_str(
     coin: &str,
     conf: &Json,
@@ -1949,16 +1951,15 @@ 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 = AddressBuilder::new(
         addr_format,
-        hash.into(),
         utxo_conf.checksum_type,
         utxo_conf.address_prefixes,
         utxo_conf.bech32_hrp,
     )
-    .as_pkh()
+    .as_pkh_from_pk(&pubkey)
     .build()?;
     address.display_address()
 }

mm2src/coins/utxo/qtum.rs

@@ -149,12 +149,11 @@ pub trait QtumBasedCoin: UtxoCommonOps + MarketCoinOps {
         let utxo = self.as_ref();
         AddressBuilder::new(
             self.addr_format().clone(),
-            AddressHashEnum::AddressHash(address.0.into()),
             utxo.conf.checksum_type,
             utxo.conf.address_prefixes.clone(),
             utxo.conf.bech32_hrp.clone(),
         )
-        .as_pkh()
+        .as_pkh(AddressHashEnum::AddressHash(address.0.into()))
         .build()
         .expect("valid address props")
     }
@@ -164,20 +163,6 @@ pub trait QtumBasedCoin: UtxoCommonOps + MarketCoinOps {
         contract_addr_from_utxo_addr(my_address).mm_err(Qrc20AddressError::from)
     }
 
-    fn utxo_address_from_contract_addr(&self, address: H160) -> Address {
-        let utxo = self.as_ref();
-        AddressBuilder::new(
-            self.addr_format().clone(),
-            AddressHashEnum::AddressHash(address.0.into()),
-            utxo.conf.checksum_type,
-            utxo.conf.address_prefixes.clone(),
-            utxo.conf.bech32_hrp.clone(),
-        )
-        .as_pkh()
-        .build()
-        .expect("valid address props")
-    }
-
     fn contract_address_from_raw_pubkey(&self, pubkey: &[u8]) -> Result<H160, String> {
         let utxo = self.as_ref();
         let qtum_address = try_s!(utxo_common::address_from_raw_pubkey(

mm2src/coins/utxo/qtum_delegation.rs

@@ -186,7 +186,7 @@ impl QtumCoin {
                         .map(|padded_staker_address_hex| padded_staker_address_hex.trim_start_matches('0'))
                 }) {
                 let hash = H160::from_str(raw).map_to_mm(|e| StakingInfosError::Internal(e.to_string()))?;
-                let address = self.utxo_address_from_contract_addr(hash);
+                let address = self.utxo_addr_from_contract_addr(hash);
                 Ok(Some(address.to_string()))
             } else {
                 Ok(None)

mm2src/coins/utxo/rpc_clients.rs

@@ -2,8 +2,8 @@
 #![cfg_attr(target_arch = "wasm32", allow(dead_code))]
 
 use crate::utxo::utxo_block_header_storage::BlockHeaderStorage;
-use crate::utxo::{output_script, sat_from_big_decimal, GetBlockHeaderError, GetConfirmedTxError, GetTxError,
-                  GetTxHeightError, ScripthashNotification};
+use crate::utxo::{output_script, output_script_p2pk, sat_from_big_decimal, GetBlockHeaderError, GetConfirmedTxError,
+                  GetTxError, GetTxHeightError, ScripthashNotification};
 use crate::{big_decimal_from_sat_unsigned, NumConversError, RpcTransportEventHandler, RpcTransportEventHandlerShared};
 use async_trait::async_trait;
 use chain::{BlockHeader, BlockHeaderBits, BlockHeaderNonce, OutPoint, Transaction as UtxoTx};
@@ -2329,12 +2329,30 @@ impl UtxoRpcClientOps for ElectrumClient {
             rpc_req!(self, "blockchain.scripthash.get_balance").into(),
             JsonRpcErrorType::Internal(err.to_string())
         )));
-        let hash = electrum_script_hash(&output_script);
-        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 hash = hex::encode(electrum_script_hash(&output_script));
+
+        // If the plain pubkey is available, fetch the balance found in P2PK output as well (if any).
+        if let Some(pubkey) = address.pubkey() {
+            let p2pk_hash = hex::encode(electrum_script_hash(&output_script_p2pk(pubkey)));
+
+            let this = self.clone();
+            let fut = async move {
+                Ok(this
+                    .scripthash_get_balances(vec![hash, p2pk_hash])
+                    .compat()
+                    .await?
+                    .into_iter()
+                    .fold(BigDecimal::from(0), |sum, electrum_balance| {
+                        sum + electrum_balance.to_big_decimal(decimals)
+                    }))
+            };
+            Box::new(fut.boxed().compat())
+        } else {
+            Box::new(
+                self.scripthash_get_balance(&hash)
+                    .map(move |electrum_balance| electrum_balance.to_big_decimal(decimals)),
+            )
+        }
     }
 
     fn display_balances(&self, addresses: Vec<Address>, decimals: u8) -> UtxoRpcFut<Vec<(Address, BigDecimal)>> {

mm2src/coins/utxo/utxo_builder/utxo_coin_builder.rs

@@ -235,12 +235,11 @@ where
     let addr_format = builder.address_format()?;
     let my_address = AddressBuilder::new(
         addr_format,
-        AddressHashEnum::AddressHash(key_pair.public().address_hash()),
         conf.checksum_type,
         conf.address_prefixes.clone(),
         conf.bech32_hrp.clone(),
     )
-    .as_pkh()
+    .as_pkh_from_pk(key_pair.public())
     .build()
     .map_to_mm(UtxoCoinBuildError::Internal)?;
 

mm2src/coins/utxo/utxo_common.rs

@@ -632,19 +632,18 @@ pub fn addresses_from_script<T: UtxoCommonOps>(coin: &T, script: &Script) -> Res
             let (addr_format, build_option) = match dst.kind {
                 AddressScriptType::P2PKH => (
                     coin.addr_format_for_standard_scripts(),
-                    AddressBuilderOption::BuildAsPubkeyHash,
+                    AddressBuilderOption::PubkeyHash(dst.hash),
                 ),
                 AddressScriptType::P2SH => (
                     coin.addr_format_for_standard_scripts(),
-                    AddressBuilderOption::BuildAsScriptHash,
+                    AddressBuilderOption::ScriptHash(dst.hash),
                 ),
-                AddressScriptType::P2WPKH => (UtxoAddressFormat::Segwit, AddressBuilderOption::BuildAsPubkeyHash),
-                AddressScriptType::P2WSH => (UtxoAddressFormat::Segwit, AddressBuilderOption::BuildAsScriptHash),
+                AddressScriptType::P2WPKH => (UtxoAddressFormat::Segwit, AddressBuilderOption::PubkeyHash(dst.hash)),
+                AddressScriptType::P2WSH => (UtxoAddressFormat::Segwit, AddressBuilderOption::ScriptHash(dst.hash)),
             };
 
             AddressBuilder::new(
                 addr_format,
-                dst.hash,
                 conf.checksum_type,
                 conf.address_prefixes.clone(),
                 conf.bech32_hrp.clone(),
@@ -1534,12 +1533,11 @@ pub async fn sign_and_send_taker_funding_spend<T: UtxoCommonOps>(
         );
         let payment_address = AddressBuilder::new(
             UtxoAddressFormat::Standard,
-            AddressHashEnum::AddressHash(dhash160(&payment_redeem_script)),
             coin.as_ref().conf.checksum_type,
             coin.as_ref().conf.address_prefixes.clone(),
             coin.as_ref().conf.bech32_hrp.clone(),
         )
-        .as_sh()
+        .as_sh(dhash160(&payment_redeem_script).into())
         .build()
         .map_err(TransactionErr::Plain)?;
         let payment_address_str = payment_address.to_string();
@@ -2732,12 +2730,11 @@ pub fn check_if_my_payment_sent<T: UtxoCommonOps + SwapOps>(
             UtxoRpcClientEnum::Native(client) => {
                 let target_addr = AddressBuilder::new(
                     coin.addr_format_for_standard_scripts(),
-                    hash.into(),
                     coin.as_ref().conf.checksum_type,
                     coin.as_ref().conf.address_prefixes.clone(),
                     coin.as_ref().conf.bech32_hrp.clone(),
                 )
-                .as_sh()
+                .as_sh(hash.into())
                 .build()?;
                 let target_addr = target_addr.to_string();
                 let is_imported = try_s!(client.is_address_imported(&target_addr).await);
@@ -4580,15 +4577,9 @@ pub fn address_from_raw_pubkey(
     hrp: Option<String>,
     addr_format: UtxoAddressFormat,
 ) -> Result<Address, String> {
-    AddressBuilder::new(
-        addr_format,
-        try_s!(Public::from_slice(pub_key)).address_hash().into(),
-        checksum_type,
-        prefixes,
-        hrp,
-    )
-    .as_pkh()
-    .build()
+    AddressBuilder::new(addr_format, checksum_type, prefixes, hrp)
+        .as_pkh_from_pk(&try_s!(Public::from_slice(pub_key)))
+        .build()
 }
 
 pub fn address_from_pubkey(
@@ -4598,8 +4589,8 @@ pub fn address_from_pubkey(
     hrp: Option<String>,
     addr_format: UtxoAddressFormat,
 ) -> Address {
-    AddressBuilder::new(addr_format, pub_key.address_hash().into(), checksum_type, prefixes, hrp)
-        .as_pkh()
+    AddressBuilder::new(addr_format, checksum_type, prefixes, hrp)
+        .as_pkh_from_pk(pub_key)
         .build()
         .expect("valid address props")
 }
@@ -4816,12 +4807,11 @@ where
 
     let payment_address = AddressBuilder::new(
         UtxoAddressFormat::Standard,
-        redeem_script_hash.into(),
         coin.as_ref().conf.checksum_type,
         coin.as_ref().conf.address_prefixes.clone(),
         coin.as_ref().conf.bech32_hrp.clone(),
     )
-    .as_sh()
+    .as_sh(redeem_script_hash.into())
     .build()?;
     let result = SwapPaymentOutputsResult {
         payment_address,

mm2src/coins/utxo/utxo_common_tests.rs

@@ -81,16 +81,10 @@ pub(super) fn utxo_coin_fields_for_test(
     } else {
         UtxoAddressFormat::Standard
     };
-    let my_address = AddressBuilder::new(
-        addr_format,
-        key_pair.public().address_hash().into(),
-        checksum_type,
-        prefixes,
-        hrp,
-    )
-    .as_pkh()
-    .build()
-    .expect("valid address props");
+    let my_address = AddressBuilder::new(addr_format, checksum_type, prefixes, hrp)
+        .as_pkh_from_pk(key_pair.public())
+        .build()
+        .expect("valid address props");
     let my_script_pubkey = Builder::build_p2pkh(my_address.hash()).to_bytes();
 
     let priv_key_policy = PrivKeyPolicy::Iguana(key_pair);

mm2src/coins/utxo/utxo_tests.rs

@@ -3088,12 +3088,11 @@ fn test_withdraw_to_p2pkh() {
     // Create a p2pkh address for the test coin
     let p2pkh_address = AddressBuilder::new(
         UtxoAddressFormat::Standard,
-        coin.as_ref().derivation_method.unwrap_single_addr().hash().clone(),
         *coin.as_ref().derivation_method.unwrap_single_addr().checksum_type(),
         coin.as_ref().conf.address_prefixes.clone(),
         coin.as_ref().conf.bech32_hrp.clone(),
     )
-    .as_pkh()
+    .as_pkh(coin.as_ref().derivation_method.unwrap_single_addr().hash().clone())
     .build()
     .expect("valid address props");
 
@@ -3138,12 +3137,11 @@ fn test_withdraw_to_p2sh() {
     // Create a p2sh address for the test coin
     let p2sh_address = AddressBuilder::new(
         UtxoAddressFormat::Standard,
-        coin.as_ref().derivation_method.unwrap_single_addr().hash().clone(),
         *coin.as_ref().derivation_method.unwrap_single_addr().checksum_type(),
         coin.as_ref().conf.address_prefixes.clone(),
         coin.as_ref().conf.bech32_hrp.clone(),
     )
-    .as_sh()
+    .as_sh(coin.as_ref().derivation_method.unwrap_single_addr().hash().clone())
     .build()
     .expect("valid address props");
 
@@ -3188,12 +3186,11 @@ fn test_withdraw_to_p2wpkh() {
     // Create a p2wpkh address for the test coin
     let p2wpkh_address = AddressBuilder::new(
         UtxoAddressFormat::Segwit,
-        coin.as_ref().derivation_method.unwrap_single_addr().hash().clone(),
         *coin.as_ref().derivation_method.unwrap_single_addr().checksum_type(),
         NetworkAddressPrefixes::default(),
         coin.as_ref().conf.bech32_hrp.clone(),
     )
-    .as_pkh()
+    .as_pkh(coin.as_ref().derivation_method.unwrap_single_addr().hash().clone())
     .build()
     .expect("valid address props");
 

mm2src/coins/utxo/utxo_withdraw.rs

@@ -13,7 +13,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};
+use keys::{AddressFormat, KeyPair, Private, Public as PublicKey};
 use mm2_core::mm_ctx::MmArc;
 use mm2_err_handle::prelude::*;
 use rpc::v1::types::ToTxHash;
@@ -469,12 +469,11 @@ where
                     .clone();
                 let my_address = AddressBuilder::new(
                     addr_format,
-                    AddressHashEnum::AddressHash(key_pair.public().address_hash()),
                     coin.as_ref().conf.checksum_type,
                     coin.as_ref().conf.address_prefixes.clone(),
                     coin.as_ref().conf.bech32_hrp.clone(),
                 )
-                .as_pkh()
+                .as_pkh_from_pk(key_pair.public())
                 .build()
                 .map_to_mm(WithdrawError::InternalError)?;
                 (key_pair, my_address)

mm2src/coins/z_coin/z_htlc.rs

@@ -48,12 +48,11 @@ pub async fn z_send_htlc(
     let script_hash = dhash160(&payment_script);
     let htlc_address = AddressBuilder::new(
         UtxoAddressFormat::Standard,
-        script_hash.into(),
         coin.utxo_arc.conf.checksum_type,
         coin.utxo_arc.conf.address_prefixes.clone(),
         None,
     )
-    .as_sh()
+    .as_sh(script_hash.into())
     .build()
     .map_to_mm(SendOutputsErr::InternalError)?;