Skip to content

Commit

Permalink
Fix for review comment
Browse files Browse the repository at this point in the history
  • Loading branch information
@Yamaguchi
Yamaguchi committed Aug 10, 2024
1 parent b33c8c9 commit 05cc6b7
Show file tree
Hide file tree
Showing 6 changed files with 371 additions and 75 deletions.
9 changes: 7 additions & 2 deletions crates/chain/src/keychain/txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@ use core::{
ops::{Bound, RangeBounds},
};
use tapyrus::{
hashes::Hash, script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, Script,
SignedAmount, Transaction, TxOut,
hashes::Hash, script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint,
PublicKey, Script, ScriptBuf, SignedAmount, Transaction, TxOut,
};

use crate::Append;
Expand Down Expand Up @@ -378,6 +378,11 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
Some((keychain.clone(), *last_index))
}

/// Insert payment base key for pay-to-contract script pubkey
pub fn insert_p2c_spk(&mut self, spk: ScriptBuf, payment_base: PublicKey) {
self.inner.insert_p2c_spk(spk, payment_base);
}

/// Returns whether the spk under the `keychain`'s `index` has been used.
///
/// Here, "unused" means that after the script pubkey was stored in the index, the index has
Expand Down
38 changes: 29 additions & 9 deletions crates/chain/src/spk_txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ use crate::{
indexed_tx_graph::Indexer,
};
use tapyrus::{
script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, Script, ScriptBuf,
SignedAmount, Transaction, TxOut,
script::color_identifier::ColorIdentifier, Amount, MalFixTxid, OutPoint, PublicKey, Script,
ScriptBuf, SignedAmount, Transaction, TxOut,
};

/// An index storing [`TxOut`]s that have a script pubkey that matches those in a list.
Expand Down Expand Up @@ -41,6 +41,8 @@ pub struct SpkTxOutIndex<I> {
txouts: BTreeMap<OutPoint, (I, TxOut)>,
/// Lookup from spk index to outpoints that had that spk
spk_txouts: BTreeSet<(I, OutPoint)>,
/// Pay-to-contract payment_base lookup by p2c spk
p2c_spks: HashMap<ScriptBuf, ScriptBuf>,
}

impl<I> Default for SpkTxOutIndex<I> {
Expand All @@ -51,6 +53,7 @@ impl<I> Default for SpkTxOutIndex<I> {
spk_indices: Default::default(),
spk_txouts: Default::default(),
unused: Default::default(),
p2c_spks: Default::default(),
}
}
}
Expand Down Expand Up @@ -103,12 +106,17 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
/// Scan a single `TxOut` for a matching script pubkey and returns the index that matches the
/// script pubkey (if any).
pub fn scan_txout(&mut self, op: OutPoint, txout: &TxOut) -> Option<&I> {
let spk_i = if txout.script_pubkey.is_colored() {
self.spk_indices.get(&ScriptBuf::from_bytes(
txout.script_pubkey.as_bytes()[35..].to_vec(),
))
let script_pubkey = if txout.script_pubkey.is_colored() {
txout.script_pubkey.remove_color()
} else {
self.spk_indices.get(&txout.script_pubkey)
txout.script_pubkey.clone()
};
let payment_base = self.p2c_spks.get(&script_pubkey);

let spk_i = if let Some(p) = payment_base {
self.spk_indices.get(p.as_script())
} else {
self.spk_indices.get(&script_pubkey)
};
if let Some(spk_i) = spk_i {
self.txouts.insert(op, (spk_i.clone(), txout.clone()));
Expand Down Expand Up @@ -192,6 +200,12 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
&self.spks
}

/// Insert payment base key for pay-to-contract script pubkey
pub fn insert_p2c_spk(&mut self, spk: ScriptBuf, payment_base: PublicKey) {
let p2c_spk = ScriptBuf::new_p2pkh(&payment_base.pubkey_hash());
self.p2c_spks.insert(spk, p2c_spk);
}

/// Adds a script pubkey to scan for. Returns `false` and does nothing if spk already exists in the map
///
/// the index will look for outputs spending to this spk whenever it scans new data.
Expand Down Expand Up @@ -304,11 +318,17 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
}
for txout in &tx.output {
let script_pubkey = if txout.script_pubkey.is_colored() {
ScriptBuf::from_bytes(txout.script_pubkey.as_bytes()[35..].to_vec())
txout.script_pubkey.remove_color()
} else {
txout.script_pubkey.clone()
};
if let Some(index) = self.index_of_spk(&script_pubkey) {
let payment_base = self.p2c_spks.get(&script_pubkey);
let script_pubkey_ref = if let Some(p) = payment_base {
p
} else {
&script_pubkey
};
if let Some(index) = self.index_of_spk(script_pubkey_ref) {
if range.contains(index)
&& txout.script_pubkey.color_id().unwrap_or_default() == *color_id
{
Expand Down
22 changes: 12 additions & 10 deletions crates/sqlite/src/store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -502,7 +502,7 @@ impl<K, A> Store<K, A> {
let payment_base: Vec<u8> = c.payment_base.to_bytes();
let spendable: u32 = if c.spendable { 1 } else { 0 };
insert_contract_stmt.execute(named_params! {
":contract_id": contract_id, ":contract": contract, ":payment_base": payment_base, ":spendable": spendable})
":contract_id": contract_id, ":contract": contract, ":payment_base": payment_base, ":spendable": spendable })
.map_err(Error::Sqlite)?;
}
Ok(())
Expand Down Expand Up @@ -663,17 +663,17 @@ mod test {
agg_changeset.unwrap().contract.get("id").unwrap().spendable,
true
);

let payment_base = PublicKey::from_str(
"028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
)
.unwrap();
let mut contract: contract::ChangeSet = contract::ChangeSet::new();
contract.insert(
"id".to_string(),
Contract {
contract_id: "id".to_string(),
contract: vec![0x00, 0x01, 0x02],
payment_base: PublicKey::from_str(
"028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
)
.unwrap(),
payment_base,
spendable: false,
},
);
Expand Down Expand Up @@ -881,6 +881,11 @@ mod test {
indexer: keychain::ChangeSet::default(),
};

let payment_base = PublicKey::from_str(
"028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
)
.unwrap();

let mut contract: contract::ChangeSet = contract::ChangeSet::new();
contract.insert(
"id".to_string(),
Expand All @@ -890,10 +895,7 @@ mod test {
0x00, 0x00, 0x55, 0x0e, 0x84, 0x00, 0xe2, 0x9b, 0x41, 0xd4, 0xa7, 0x16, 0x44,
0x66, 0x55, 0x44, 0x00, 0x00,
],
payment_base: PublicKey::from_str(
"028bde91b10013e08949a318018fedbd896534a549a278e220169ee2a36517c7aa",
)
.unwrap(),
payment_base,
spendable: true,
},
);
Expand Down
116 changes: 81 additions & 35 deletions crates/wallet/src/wallet/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,11 @@ pub enum GenerateContractError {
},
/// Invalid address
InvalidAddress(tapyrus::address::Error),
/// Payment base is invalid.
InvalidPublicKey {
/// public key(payment base).
public_key: PublicKey,
},
}

impl fmt::Display for GenerateContractError {
Expand All @@ -448,6 +453,9 @@ impl fmt::Display for GenerateContractError {
GenerateContractError::ContractTooLarge { length } => {
write!(f, "contract is too large (size: {})", length)
}
GenerateContractError::InvalidPublicKey { public_key } => {
write!(f, "invalid payment base ({})", public_key)
}
GenerateContractError::InvalidAddress(e) => e.fmt(f),
}
}
Expand All @@ -464,6 +472,11 @@ pub enum CreateContractError {
/// identifier of contract.
contract_id: String,
},
/// Payment base is invalid.
InvalidPublicKey {
/// public key(payment base).
public_key: PublicKey,
},
/// Other error.
Error {
/// An error that caused this error.
Expand All @@ -477,6 +490,9 @@ impl fmt::Display for CreateContractError {
CreateContractError::ContractAlreadyExist { contract_id } => {
write!(f, "contract already exists (contract_id: {})", contract_id)
}
CreateContractError::InvalidPublicKey { public_key } => {
write!(f, "invalid payment base ({})", public_key)
}
CreateContractError::Error { reason } => {
write!(f, "can not create contract address (reason: {})", reason)
}
Expand Down Expand Up @@ -546,6 +562,18 @@ impl Wallet {
)
}

fn build_p2c_spk_from_contract(wallet: &mut Wallet, contracts: contract::ChangeSet) {
for (_, contract) in contracts {
let payment_base = contract.payment_base;
match wallet.create_pay_to_contract_script(&payment_base, contract.contract, None) {
Ok(spk) => {
wallet.insert_p2c_spk(spk, payment_base);
}
_ => {}
}
}
}

/// Initialize an empty [`Wallet`] with a custom genesis hash.
///
/// This is like [`Wallet::new`] with an additional `genesis_hash` parameter. This is useful
Expand All @@ -572,7 +600,7 @@ impl Wallet {

let indexed_graph = IndexedTxGraph::new(index);

let contracts = contract::ChangeSet::new();
let contracts: BTreeMap<String, Contract> = contract::ChangeSet::new();

let mut persist = Persist::new(db);
persist.stage(ChangeSet {
Expand All @@ -583,16 +611,18 @@ impl Wallet {
});
persist.commit().map_err(NewError::Persist)?;

Ok(Wallet {
let mut wallet = Wallet {
signers,
change_signers,
network,
chain,
indexed_graph,
persist,
secp,
contracts,
})
contracts: contracts.clone(),
};

Ok(wallet)
}

/// Load [`Wallet`] from the given persistence backend.
Expand Down Expand Up @@ -682,16 +712,18 @@ impl Wallet {

let persist = Persist::new(db);

Ok(Wallet {
let mut wallet = Wallet {
signers,
change_signers,
chain,
indexed_graph,
persist,
network,
secp,
contracts,
})
contracts: contracts.clone(),
};
Self::build_p2c_spk_from_contract(&mut wallet, contracts);
Ok(wallet)
}

/// Either loads [`Wallet`] from persistence, or initializes it if it does not exist.
Expand Down Expand Up @@ -982,17 +1014,26 @@ impl Wallet {
contract: Vec<u8>,
color_id: Option<ColorIdentifier>,
) -> Result<Address<NetworkChecked>, GenerateContractError> {
let script = self.create_pay_to_contract_script(payment_base, contract, color_id)?;
let address = Address::from_script(script.as_script(), self.network)
.map_err(GenerateContractError::InvalidAddress)?;
Ok(address)
}

fn create_pay_to_contract_script(
&self,
payment_base: &PublicKey,
contract: Vec<u8>,
color_id: Option<ColorIdentifier>,
) -> Result<ScriptBuf, GenerateContractError> {
let p2c_public_key = self.pay_to_contract_key(payment_base, contract)?;
let pubkey_hash = p2c_public_key.pubkey_hash();
let script: ScriptBuf = match color_id {
Some(c) if c.is_colored() => ScriptBuf::new_cp2pkh(&color_id.unwrap(), &pubkey_hash),
_ => ScriptBuf::new_p2pkh(&pubkey_hash),
};
let address = Address::from_script(script.as_script(), self.network)
.map_err(GenerateContractError::InvalidAddress)?;
Ok(address)
Ok(script)
}

/// Generate pay-to-contract public key with the specified content hash.
pub fn pay_to_contract_key(
&self,
Expand All @@ -1007,6 +1048,18 @@ impl Wallet {
length: contract.len(),
});
}
let index = match self
.indexed_graph
.index
.index_of_spk(ScriptBuf::new_p2pkh(&payment_base.pubkey_hash()).as_script())
{
Some(i) => i,
None => {
return Err(GenerateContractError::InvalidPublicKey {
public_key: payment_base.clone(),
})
}
};
let commitment: Scalar = self.create_pay_to_contract_commitment(payment_base, contract);
let pubkey = payment_base
.inner
Expand Down Expand Up @@ -2683,15 +2736,8 @@ impl Wallet {
&self.indexed_graph.index
}

/// Insert a descriptor with a keychain associated to it.
pub fn insert_descriptor(
&mut self,
keychain: KeychainKind,
descriptor: Descriptor<DescriptorPublicKey>,
) -> tdk_chain::keychain::txout_index::ChangeSet<KeychainKind> {
self.indexed_graph
.index
.insert_descriptor(keychain, descriptor)
fn insert_p2c_spk(&mut self, spk: ScriptBuf, payment_base: PublicKey) {
self.indexed_graph.index.insert_p2c_spk(spk, payment_base)
}

/// Get a reference to the inner [`LocalChain`].
Expand Down Expand Up @@ -2778,11 +2824,18 @@ impl Wallet {
contract: Vec<u8>,
payment_base: PublicKey,
spendable: bool,
) -> Result<(), CreateContractError> {
) -> Result<Contract, CreateContractError> {
let mut changeset = ChangeSet::default();
if let Some(c) = self.contracts.get(&contract_id) {
return Err(CreateContractError::ContractAlreadyExist { contract_id });
} else {
}
let new_contract = {
let spk = self
.create_pay_to_contract_script(&payment_base, contract.clone(), None)
.map_err(|e| CreateContractError::InvalidPublicKey {
public_key: payment_base,
})?;
self.insert_p2c_spk(spk.clone(), payment_base);
let new_contract = Contract {
contract_id: contract_id.clone(),
contract: contract.clone(),
Expand All @@ -2792,24 +2845,17 @@ impl Wallet {
changeset
.contract
.insert(contract_id.clone(), new_contract.clone());
let p2c_public_key =
self.pay_to_contract_key(&payment_base, contract)
.map_err(|e| CreateContractError::Error {
reason: e.to_string(),
})?;
let descriptor_str = format!("pkh({})", p2c_public_key);
let (descriptor, _) =
Descriptor::<DescriptorPublicKey>::parse_descriptor(&self.secp, &descriptor_str)
.unwrap();
let _ = self.insert_descriptor(KeychainKind::External, descriptor);
self.contracts.insert(contract_id.clone(), new_contract);

self.contracts
.insert(contract_id.clone(), new_contract.clone());
self.persist
.stage_and_commit(changeset)
.map_err(|e| CreateContractError::Error {
reason: e.to_string(),
})?;
}
Ok(())
new_contract
};
Ok(new_contract)
}

/// Update pay-to-contract information to the wallet.
Expand Down
Loading

0 comments on commit 05cc6b7

Please sign in to comment.