1use blst::min_pk::{AggregateSignature, PublicKey, SecretKey, Signature};
2use dg_xch_core::blockchain::coin_spend::CoinSpend;
3use dg_xch_core::blockchain::condition_with_args::Message;
4use dg_xch_core::blockchain::sized_bytes::{Bytes32, Bytes48, Bytes96};
5use dg_xch_core::blockchain::spend_bundle::SpendBundle;
6use dg_xch_core::blockchain::unsized_bytes::UnsizedBytes;
7use dg_xch_core::blockchain::utils::pkm_pairs_for_conditions;
8use dg_xch_core::blockchain::wallet_type::WalletType;
9use dg_xch_core::clvm::bls_bindings;
10use dg_xch_core::clvm::bls_bindings::{aggregate_verify_signature, verify_signature};
11use dg_xch_core::clvm::condition_utils::conditions_for_solution;
12use dg_xch_core::consensus::constants::ConsensusConstants;
13use dg_xch_core::traits::SizedBytes;
14use log::{debug, error, info, warn};
15use num_traits::cast::ToPrimitive;
16use std::collections::HashMap;
17use std::future::Future;
18use std::io::Error;
19
20pub struct DerivationRecord {
21 pub index: u32,
22 pub puzzle_hash: Bytes32,
23 pub pubkey: Bytes48,
24 pub wallet_type: WalletType,
25 pub wallet_id: u32,
26 pub hardened: bool,
27}
28
29pub async fn sign_coin_spend<F, Fut>(
30 coin_spend: CoinSpend,
31 key_fn: F,
32 pre_calculated_signatures: HashMap<(Bytes48, Message), Bytes96>,
33 constants: &ConsensusConstants,
34) -> Result<SpendBundle, Error>
35where
36 F: Fn(&Bytes48) -> Fut,
37 Fut: Future<Output = Result<SecretKey, Error>>,
38{
39 sign_coin_spends(
40 vec![coin_spend],
41 key_fn,
42 pre_calculated_signatures,
43 &constants.agg_sig_me_additional_data,
44 constants.max_block_cost_clvm.to_u64().unwrap(),
45 )
46 .await
47}
48
49pub async fn sign_coin_spends<F, Fut>(
50 coin_spends: Vec<CoinSpend>,
51 key_fn: F,
52 pre_calculated_signatures: HashMap<(Bytes48, Message), Bytes96>,
53 additional_data: &[u8],
54 max_cost: u64,
55) -> Result<SpendBundle, Error>
56where
57 F: Fn(&Bytes48) -> Fut,
58 Fut: Future<Output = Result<SecretKey, Error>>,
59{
60 let mut signatures: Vec<Signature> = vec![];
61 let mut pk_list: Vec<Bytes48> = vec![];
62 let mut msg_list: Vec<Vec<u8>> = vec![];
63 debug!("Creating Signatures for Coin Spends");
64 for coin_spend in &coin_spends {
65 let conditions =
67 conditions_for_solution(&coin_spend.puzzle_reveal, &coin_spend.solution, max_cost)?.0;
68 for (code, pk_bytes, msg) in
70 pkm_pairs_for_conditions(&conditions, coin_spend.coin, additional_data)?
71 {
72 let pk = PublicKey::from_bytes(pk_bytes.as_ref()).map_err(|e| {
73 Error::other(format!(
74 "Failed to parse Public key: {}, {:?}",
75 hex::encode(pk_bytes),
76 e
77 ))
78 })?;
79 let secret_key = (key_fn)(&pk_bytes).await?;
80 let signature = if secret_key.sk_to_pk() != pk {
81 debug!("Searching for PreCalc on op {code} ({pk_bytes}, {msg})");
82 pre_calculated_signatures.get(&(pk_bytes, msg)).ok_or_else(|| {
84 error!("Failed to find ({pk_bytes}, {msg}) in map \n {pre_calculated_signatures:#?}");
85 Error::other(
86 format!(
87 "Failed to find Secret Key for Public Key: {}",
88 Bytes48::new(pk.to_bytes())
89 ),
90 )
91 })?.try_into()?
92 } else {
93 assert_eq!(&secret_key.sk_to_pk(), &pk);
94 bls_bindings::sign(&secret_key, msg.as_ref())
95 };
96 if !verify_signature(&pk, msg.as_ref(), &signature) {
97 return Err(Error::other(format!(
98 "Failed to find Validate Signature for Message: {} - {}",
99 code,
100 UnsizedBytes::new(msg.as_ref())
101 )));
102 }
103 pk_list.push(pk_bytes);
104 msg_list.push(msg.as_ref().to_vec());
105 signatures.push(signature);
106 }
107 }
108 debug!("Creating Aggregate signature");
109 let sig_refs: Vec<&Signature> = signatures.iter().collect();
110 let msg_list: Vec<&[u8]> = msg_list.iter().map(Vec::as_slice).collect();
111 let aggsig = AggregateSignature::aggregate(&sig_refs, true)
112 .map_err(|e| Error::other(format!("Failed to aggregate signatures: {e:?}")))?
113 .to_signature();
114 assert!(aggregate_verify_signature(&pk_list, &msg_list, &aggsig));
115 Ok(SpendBundle {
116 coin_spends,
117 aggregated_signature: Bytes96::from(aggsig),
118 })
119}
120
121pub async fn partial_signature<F, Fut>(
122 coin_spends: Vec<CoinSpend>,
123 key_fn: F,
124 pre_calculated_signatures: HashMap<(Bytes48, Message), Bytes96>,
125 additional_data: &[u8],
126 max_cost: u64,
127) -> Result<SpendBundle, Error>
128where
129 F: Fn(&Bytes48) -> Fut,
130 Fut: Future<Output = Result<SecretKey, Error>>,
131{
132 let mut signatures: Vec<Signature> = vec![];
133 let mut pk_list: Vec<Bytes48> = vec![];
134 let mut msg_list: Vec<Vec<u8>> = vec![];
135 debug!("Creating Signatures for Coin Spends");
136 for coin_spend in &coin_spends {
137 let conditions_dict =
139 conditions_for_solution(&coin_spend.puzzle_reveal, &coin_spend.solution, max_cost)?.0;
140 let mut total_messages = 0;
142 let mut signed_messages = 0;
143 for (code, pk_bytes, msg) in
144 pkm_pairs_for_conditions(&conditions_dict, coin_spend.coin, additional_data)?
145 {
146 let pk = PublicKey::from_bytes(pk_bytes.as_ref()).map_err(|e| {
147 Error::other(format!(
148 "Failed to parse Public key: {}, {:?}",
149 hex::encode(pk_bytes),
150 e
151 ))
152 })?;
153 total_messages += 1;
154 if let Ok(secret_key) = (key_fn)(&pk_bytes).await {
155 let signature = if secret_key.sk_to_pk() != pk {
156 if let Some(signature) = pre_calculated_signatures.get(&(pk_bytes, msg)) {
157 Some(signature.try_into()?)
158 } else {
159 None
160 }
161 } else {
162 Some(bls_bindings::sign(&secret_key, msg.as_ref()))
163 };
164 if let Some(signature) = signature {
165 info!("Signing Partial Message");
166 if !verify_signature(&pk, msg.as_ref(), &signature) {
167 return Err(Error::other(format!(
168 "Failed to find Validate Signature for Message: {code} - {}",
169 UnsizedBytes::new(msg.as_ref())
170 )));
171 }
172 signed_messages += 1;
173 pk_list.push(pk_bytes);
174 msg_list.push(msg.as_ref().to_vec());
175 signatures.push(signature);
176 } else {
177 warn!("Got Secret Key but No Signature for Partial Message");
178 }
179 }
180 }
181 info!("Signed {signed_messages}/{total_messages} messages");
182 }
183 let spend_bundle = if !signatures.is_empty() {
184 info!("Creating Aggregate signature");
185 let sig_refs: Vec<&Signature> = signatures.iter().collect();
186 let msg_list: Vec<&[u8]> = msg_list.iter().map(Vec::as_slice).collect();
187 let aggsig = AggregateSignature::aggregate(&sig_refs, true)
188 .map_err(|e| Error::other(format!("Failed to aggregate signatures: {e:?}")))?
189 .to_signature();
190 if !aggregate_verify_signature(&pk_list, &msg_list, &aggsig) {
191 return Err(Error::other("Failed to Validate Aggregate Signature"));
192 }
193 assert!(aggregate_verify_signature(&pk_list, &msg_list, &aggsig));
194 SpendBundle {
195 coin_spends,
196 aggregated_signature: Bytes96::from(aggsig),
197 }
198 } else {
199 info!("Empty Signature List");
200 SpendBundle {
201 coin_spends,
202 aggregated_signature: Bytes96::default(),
203 }
204 };
205 Ok(spend_bundle)
206}