1  use std::sync::Arc;
  2  
  3  use bitcoin::key::KeyPair;
  4  use fedimint_core::core::{DynInput, DynOutput, IntoDynInstance, ModuleInstanceId};
  5  use fedimint_core::transaction::{Transaction, TransactionSignature};
  6  use fedimint_core::Amount;
  7  use itertools::multiunzip;
  8  use rand::{CryptoRng, Rng, RngCore};
  9  use secp256k1_zkp::Secp256k1;
 10  
 11  use crate::module::StateGenerator;
 12  use crate::sm::DynState;
 13  
 14  #[derive(Clone)]
 15  pub struct ClientInput<I = DynInput, S = DynState> {
 16      pub input: I,
 17      pub keys: Vec<KeyPair>,
 18      pub amount: Amount,
 19      pub state_machines: StateGenerator<S>,
 20  }
 21  
 22  impl<I, S> IntoDynInstance for ClientInput<I, S>
 23  where
 24      I: IntoDynInstance<DynType = DynInput> + 'static,
 25      S: IntoDynInstance<DynType = DynState> + 'static,
 26  {
 27      type DynType = ClientInput;
 28  
 29      fn into_dyn(self, module_instance_id: ModuleInstanceId) -> ClientInput {
 30          ClientInput {
 31              input: self.input.into_dyn(module_instance_id),
 32              keys: self.keys,
 33              amount: self.amount,
 34              state_machines: state_gen_to_dyn(self.state_machines, module_instance_id),
 35          }
 36      }
 37  }
 38  
 39  #[derive(Clone)]
 40  pub struct ClientOutput<O = DynOutput, S = DynState> {
 41      pub output: O,
 42      pub amount: Amount,
 43      pub state_machines: StateGenerator<S>,
 44  }
 45  
 46  impl<O, S> IntoDynInstance for ClientOutput<O, S>
 47  where
 48      O: IntoDynInstance<DynType = DynOutput> + 'static,
 49      S: IntoDynInstance<DynType = DynState> + 'static,
 50  {
 51      type DynType = ClientOutput;
 52  
 53      fn into_dyn(self, module_instance_id: ModuleInstanceId) -> ClientOutput {
 54          ClientOutput {
 55              output: self.output.into_dyn(module_instance_id),
 56              amount: self.amount,
 57              state_machines: state_gen_to_dyn(self.state_machines, module_instance_id),
 58          }
 59      }
 60  }
 61  
 62  #[derive(Default, Clone)]
 63  pub struct TransactionBuilder {
 64      pub(crate) inputs: Vec<ClientInput>,
 65      pub(crate) outputs: Vec<ClientOutput>,
 66  }
 67  
 68  impl TransactionBuilder {
 69      pub fn new() -> Self {
 70          Self::default()
 71      }
 72  
 73      pub fn with_input(mut self, input: ClientInput) -> Self {
 74          self.inputs.push(input);
 75          self
 76      }
 77  
 78      pub fn with_output(mut self, output: ClientOutput) -> Self {
 79          self.outputs.push(output);
 80          self
 81      }
 82  
 83      pub fn with_inputs(mut self, inputs: Vec<ClientInput>) -> Self {
 84          for input in inputs {
 85              self.inputs.push(input);
 86          }
 87  
 88          self
 89      }
 90  
 91      pub fn with_outputs(mut self, outputs: Vec<ClientOutput>) -> Self {
 92          for output in outputs {
 93              self.outputs.push(output);
 94          }
 95  
 96          self
 97      }
 98  
 99      pub fn build<C, R: RngCore + CryptoRng>(
100          self,
101          secp_ctx: &Secp256k1<C>,
102          mut rng: R,
103      ) -> (Transaction, Vec<DynState>)
104      where
105          C: secp256k1_zkp::Signing + secp256k1_zkp::Verification,
106      {
107          let (inputs, input_keys, input_states): (Vec<_>, Vec<_>, Vec<_>) = multiunzip(
108              self.inputs
109                  .into_iter()
110                  .map(|input| (input.input, input.keys, input.state_machines)),
111          );
112          let (outputs, output_states): (Vec<_>, Vec<_>) = self
113              .outputs
114              .into_iter()
115              .map(|output| (output.output, output.state_machines))
116              .unzip();
117  
118          let nonce: [u8; 8] = rng.gen();
119  
120          let txid = Transaction::tx_hash_from_parts(&inputs, &outputs, nonce);
121          let msg = secp256k1_zkp::Message::from_slice(&txid[..]).expect("txid has right length");
122  
123          let signatures = input_keys
124              .into_iter()
125              .flatten()
126              .map(|keypair| secp_ctx.sign_schnorr(&msg, &keypair))
127              .collect();
128  
129          let transaction = Transaction {
130              inputs,
131              outputs,
132              nonce,
133              signatures: TransactionSignature::NaiveMultisig(signatures),
134          };
135  
136          let states = input_states
137              .into_iter()
138              .enumerate()
139              .chain(output_states.into_iter().enumerate())
140              .flat_map(|(idx, state_gen)| state_gen(txid, idx as u64))
141              .collect::<Vec<_>>();
142  
143          (transaction, states)
144      }
145  }
146  
147  fn state_gen_to_dyn<S>(
148      state_gen: StateGenerator<S>,
149      module_instance: ModuleInstanceId,
150  ) -> StateGenerator<DynState>
151  where
152      S: IntoDynInstance<DynType = DynState> + 'static,
153  {
154      Arc::new(move |txid, index| {
155          let states = state_gen(txid, index);
156          states
157              .into_iter()
158              .map(|state| state.into_dyn(module_instance))
159              .collect()
160      })
161  }