Support refunds as possible for invalidly received outputs on Serai

processor/src/multisigs/db.rs

@@ -5,7 +5,11 @@ use ciphersuite::Ciphersuite;
 pub use serai_db::*;
 
 use scale::{Encode, Decode};
-use serai_client::in_instructions::primitives::InInstructionWithBalance;
+#[rustfmt::skip]
+use serai_client::{
+  primitives::ExternalAddress,
+  in_instructions::primitives::InInstructionWithBalance
+};
 
 use crate::{
   Get, Db, Plan,
@@ -153,6 +157,21 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
     txn.put(Self::resolved_key(resolution.as_ref()), plan);
   }
 
+  fn refund_key(id: &[u8]) -> Vec<u8> {
+    Self::multisigs_key(b"refund", id)
+  }
+  pub fn set_refund(txn: &mut D::Transaction<'_>, id: &[u8], address: ExternalAddress) {
+    txn.put(Self::refund_key(id), address.encode());
+  }
+  pub fn take_refund(txn: &mut D::Transaction<'_>, id: &[u8]) -> Option<ExternalAddress> {
+    let key = Self::refund_key(id);
+    let res = txn.get(&key).map(|address| ExternalAddress::decode(&mut address.as_ref()).unwrap());
+    if res.is_some() {
+      txn.del(key);
+    }
+    res
+  }
+
   fn forwarded_output_key(amount: u64) -> Vec<u8> {
     Self::multisigs_key(b"forwarded_output", amount.to_le_bytes())
   }

processor/src/multisigs/mod.rs

@@ -7,7 +7,7 @@ use scale::{Encode, Decode};
 use messages::SubstrateContext;
 
 use serai_client::{
-  primitives::{BlockHash, MAX_DATA_LEN},
+  primitives::{MAX_DATA_LEN, ExternalAddress, BlockHash},
   in_instructions::primitives::{
     InInstructionWithBalance, Batch, RefundableInInstruction, Shorthand, MAX_BATCH_SIZE,
   },
@@ -40,9 +40,21 @@ use crate::{
 };
 
 // InInstructionWithBalance from an external output
-fn instruction_from_output<N: Network>(output: &N::Output) -> Option<InInstructionWithBalance> {
+fn instruction_from_output<N: Network>(
+  output: &N::Output,
+) -> (Option<ExternalAddress>, Option<InInstructionWithBalance>) {
   assert_eq!(output.kind(), OutputType::External);
 
+  let presumed_origin = output.presumed_origin().map(|address| {
+    ExternalAddress::new(
+      address
+        .try_into()
+        .map_err(|_| ())
+        .expect("presumed origin couldn't be converted to a Vec<u8>"),
+    )
+    .expect("presumed origin exceeded address limits")
+  });
+
   let mut data = output.data();
   let max_data_len = usize::try_from(MAX_DATA_LEN).unwrap();
   if data.len() > max_data_len {
@@ -51,19 +63,23 @@ fn instruction_from_output<N: Network>(output: &N::Output) -> Option<InInstructi
       hex::encode(output.id()),
       data.len(),
     );
-    None?;
+    return (presumed_origin, None);
   }
 
-  let Ok(shorthand) = Shorthand::decode(&mut data) else { None? };
+  let Ok(shorthand) = Shorthand::decode(&mut data) else { return (presumed_origin, None) };
-  let Ok(instruction) = RefundableInInstruction::try_from(shorthand) else { None? };
+  let Ok(instruction) = RefundableInInstruction::try_from(shorthand) else {
+    return (presumed_origin, None);
+  };
 
   let mut balance = output.balance();
   // Deduct twice the cost to aggregate to prevent economic attacks by malicious miners against
   // other users
   balance.amount.0 -= 2 * N::COST_TO_AGGREGATE;
 
-  // TODO2: Set instruction.origin if not set (and handle refunds in general)
+  (
-  Some(InInstructionWithBalance { instruction: instruction.instruction, balance })
+    instruction.origin.or(presumed_origin),
+    Some(InInstructionWithBalance { instruction: instruction.instruction, balance }),
+  )
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
@@ -320,6 +336,17 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
     (existing_outputs, new_outputs)
   }
 
+  fn refund_plan(output: &N::Output, refund_to: N::Address) -> Plan<N> {
+    Plan {
+      key: output.key(),
+      inputs: vec![output.clone()],
+      // Uses a payment as this will still be successfully sent due to fee amortization,
+      // and because change is currently always a Serai key
+      payments: vec![Payment { address: refund_to, data: None, amount: output.balance().amount.0 }],
+      change: None,
+    }
+  }
+
   // Manually creates Plans for all External outputs needing forwarding/refunding.
   //
   // Returns created Plans and a map of forwarded output IDs to their associated InInstructions.
@@ -351,8 +378,10 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
     let mut plans = vec![];
     let mut forwarding = HashMap::new();
     existing_outputs.retain(|output| {
+      let plans_at_start = plans.len();
       if output.kind() == OutputType::External {
-        if let Some(instruction) = instruction_from_output::<N>(output) {
+        let (refund_to, instruction) = instruction_from_output::<N>(output);
+        if let Some(instruction) = instruction {
           // Build a dedicated Plan forwarding this
           plans.push(Plan {
             key: self.existing.as_ref().unwrap().key,
@@ -363,13 +392,15 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
 
           // Set the instruction for this output to be returned
           forwarding.insert(output.id().as_ref().to_vec(), instruction);
+        } else if let Some(refund_to) = refund_to {
+          if let Ok(refund_to) = refund_to.consume().try_into() {
+            // Build a dedicated Plan refunding this
+            plans.push(Self::refund_plan(output, refund_to));
           }
-
-        // TODO: Refund here
-        false
-      } else {
-        true
         }
+      }
+      // Only keep if we didn't make a Plan consuming it
+      plans_at_start == plans.len()
     });
     (plans, forwarding)
   }
@@ -571,7 +602,7 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
 
     // We now have to acknowledge the acknowledged block, if it's new
     // It won't be if this block's `InInstruction`s were split into multiple `Batch`s
-    let (acquired_lock, (mut existing_outputs, new_outputs)) = {
+    let (acquired_lock, (mut existing_outputs, mut new_outputs)) = {
       let (acquired_lock, outputs) = if ScannerHandle::<N, D>::db_scanned(txn)
         .expect("published a Batch despite never scanning a block") <
         block_number
@@ -603,6 +634,21 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
       }
     };
 
+    let handle_refund_outputs = |txn: &mut _, plans: &mut Vec<_>, outputs: &mut Vec<N::Output>| {
+      outputs.retain(|output| {
+        if let Some(refund_to) = MultisigsDb::<N, D>::take_refund(txn, output.id().as_ref()) {
+          // If this isn't a valid refund address, accumulate this output
+          let Ok(refund_to) = refund_to.consume().try_into() else {
+            return true;
+          };
+          plans.push(Self::refund_plan(output, refund_to));
+          return false;
+        }
+        true
+      });
+    };
+    handle_refund_outputs(txn, &mut plans, &mut existing_outputs);
+
     plans.extend({
       let existing = self.existing.as_mut().unwrap();
       let existing_key = existing.key;
@@ -636,6 +682,7 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
       }
     }
 
+    handle_refund_outputs(txn, &mut plans, &mut new_outputs);
     if let Some(new) = self.new.as_mut() {
       plans.extend(new.scheduler.schedule::<D>(txn, new_outputs, new_payments, new.key, false));
     }
@@ -780,11 +827,20 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
             continue;
           }
 
-          let instruction = if let Some(instruction) = instruction_from_output::<N>(&output) {
+          let (refund_to, instruction) = instruction_from_output::<N>(&output);
+          let refund = |txn| {
+            if let Some(refund_to) = refund_to {
+              MultisigsDb::<N, D>::set_refund(txn, output.id().as_ref(), refund_to);
+            }
+          };
+          let instruction = if let Some(instruction) = instruction {
             instruction
           } else {
+            // If the output data is empty, this may be a forwarded transaction from a prior
+            // multisig
+            // If it's not empty, it's corrupt in some way and should be refunded
             if !output.data().is_empty() {
-              // TODO2: Refund
+              refund(txn);
               continue;
             }
 
@@ -793,7 +849,8 @@ impl<D: Db, N: Network> MultisigManager<D, N> {
             {
               instruction
             } else {
-              // TODO2: Refund
+              // If it's not a forwarded output, refund
+              refund(txn);
               continue;
             }
           };

processor/src/networks/bitcoin.rs

@@ -1,6 +1,9 @@
 use std::{time::Duration, io, collections::HashMap};
 
 use async_trait::async_trait;
+use once_cell::sync::Lazy;
+
+use scale::{Encode, Decode};
 
 use transcript::{Transcript, RecommendedTranscript};
 use ciphersuite::group::ff::PrimeField;
@@ -19,7 +22,7 @@ use bitcoin_serai::{
     consensus::{Encodable, Decodable},
     script::Instruction,
     address::{NetworkChecked, Address as BAddress},
-    OutPoint, TxOut, Transaction, Block, Network as BitcoinNetwork,
+    OutPoint, TxOut, Transaction, Block, Network as BNetwork,
   },
   wallet::{
     tweak_keys, address_payload, ReceivedOutput, Scanner, TransactionError,
@@ -53,8 +56,6 @@ use crate::{
   Payment,
 };
 
-use once_cell::sync::Lazy;
-
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct OutputId(pub [u8; 36]);
 impl Default for OutputId {
@@ -76,6 +77,7 @@ impl AsMut<[u8]> for OutputId {
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Output {
   kind: OutputType,
+  presumed_origin: Option<Address>,
   output: ReceivedOutput,
   data: Vec<u8>,
 }
@@ -119,6 +121,10 @@ impl OutputTrait<Bitcoin> for Output {
       (ProjectivePoint::GENERATOR * self.output.offset())
   }
 
+  fn presumed_origin(&self) -> Option<Address> {
+    self.presumed_origin.clone()
+  }
+
   fn balance(&self) -> Balance {
     Balance { coin: SeraiCoin::Bitcoin, amount: Amount(self.output.value()) }
   }
@@ -129,14 +135,25 @@ impl OutputTrait<Bitcoin> for Output {
 
   fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
     self.kind.write(writer)?;
+    let presumed_origin: Option<Vec<u8>> =
+      self.presumed_origin.clone().map(|presumed_origin| presumed_origin.try_into().unwrap());
+    writer.write_all(&presumed_origin.encode())?;
     self.output.write(writer)?;
     writer.write_all(&u16::try_from(self.data.len()).unwrap().to_le_bytes())?;
     writer.write_all(&self.data)
   }
 
-  fn read<R: io::Read>(reader: &mut R) -> io::Result<Self> {
+  fn read<R: io::Read>(mut reader: &mut R) -> io::Result<Self> {
     Ok(Output {
       kind: OutputType::read(reader)?,
+      presumed_origin: {
+        let mut io_reader = scale::IoReader(reader);
+        let res = Option::<Vec<u8>>::decode(&mut io_reader)
+          .unwrap()
+          .map(|address| Address::try_from(address).unwrap());
+        reader = io_reader.0;
+        res
+      },
       output: ReceivedOutput::read(reader)?,
       data: {
         let mut data_len = [0; 2];
@@ -494,7 +511,7 @@ impl Network for Bitcoin {
   }
 
   fn address(key: ProjectivePoint) -> Address {
-    Address(BAddress::<NetworkChecked>::new(BitcoinNetwork::Bitcoin, address_payload(key).unwrap()))
+    Address(BAddress::<NetworkChecked>::new(BNetwork::Bitcoin, address_payload(key).unwrap()))
   }
 
   fn branch_address(key: ProjectivePoint) -> Address {
@@ -545,7 +562,35 @@ impl Network for Bitcoin {
         };
         data.truncate(MAX_DATA_LEN.try_into().unwrap());
 
-        let output = Output { kind, output, data };
+        let presumed_origin = {
+          let spent_output = tx.input[0].previous_output;
+          let mut spent_tx = spent_output.txid.as_raw_hash().to_byte_array();
+          spent_tx.reverse();
+          let spent_output = {
+            let mut tx;
+            while {
+              tx = self.get_transaction(&spent_tx).await;
+              tx.is_err()
+            } {
+              log::error!("couldn't get transaction from bitcoin node: {tx:?}");
+              sleep(Duration::from_secs(5)).await;
+            }
+            tx.unwrap().output.swap_remove(usize::try_from(spent_output.vout).unwrap())
+          };
+          BAddress::from_script(&spent_output.script_pubkey, BNetwork::Bitcoin).ok().and_then(
+            |address| {
+              let address = Address(address);
+              let encoded: Result<Vec<u8>, _> = address.clone().try_into();
+              if encoded.is_ok() {
+                Some(address)
+              } else {
+                None
+              }
+            },
+          )
+        };
+
+        let output = Output { kind, presumed_origin, output, data };
         assert_eq!(output.tx_id(), tx.id());
         outputs.push(output);
       }
@@ -717,7 +762,7 @@ impl Network for Bitcoin {
       .rpc
       .rpc_call::<Vec<String>>(
         "generatetoaddress",
-        serde_json::json!([1, BAddress::p2sh(Script::new(), BitcoinNetwork::Regtest).unwrap()]),
+        serde_json::json!([1, BAddress::p2sh(Script::new(), BNetwork::Regtest).unwrap()]),
       )
       .await
       .unwrap();
@@ -726,9 +771,9 @@ impl Network for Bitcoin {
   #[cfg(test)]
   async fn test_send(&self, address: Address) -> Block {
     let secret_key = SecretKey::new(&mut rand_core::OsRng);
-    let private_key = PrivateKey::new(secret_key, BitcoinNetwork::Regtest);
+    let private_key = PrivateKey::new(secret_key, BNetwork::Regtest);
     let public_key = PublicKey::from_private_key(SECP256K1, &private_key);
-    let main_addr = BAddress::p2pkh(&public_key, BitcoinNetwork::Regtest);
+    let main_addr = BAddress::p2pkh(&public_key, BNetwork::Regtest);
 
     let new_block = self.get_latest_block_number().await.unwrap() + 1;
     self

processor/src/networks/mod.rs

@@ -103,6 +103,8 @@ pub trait Output<N: Network>: Send + Sync + Sized + Clone + PartialEq + Eq + Deb
   fn tx_id(&self) -> <N::Transaction as Transaction<N>>::Id;
   fn key(&self) -> <N::Curve as Ciphersuite>::G;
 
+  fn presumed_origin(&self) -> Option<N::Address>;
+
   fn balance(&self) -> Balance;
   // TODO: Remove this?
   fn amount(&self) -> u64 {

processor/src/networks/monero.rs

@@ -77,6 +77,10 @@ impl OutputTrait<Monero> for Output {
     EdwardsPoint(self.0.output.data.key - (EdwardsPoint::generator().0 * self.0.key_offset()))
   }
 
+  fn presumed_origin(&self) -> Option<Address> {
+    None
+  }
+
   fn balance(&self) -> Balance {
     Balance { coin: SeraiCoin::Monero, amount: Amount(self.0.commitment().amount) }
   }

processor/src/plan.rs

@@ -10,6 +10,7 @@ use crate::networks::{Output, Network};
 pub struct Payment<N: Network> {
   pub address: N::Address,
   pub data: Option<Vec<u8>>,
+  // TODO: Balance
   pub amount: u64,
 }
 

substrate/client/src/networks/bitcoin.rs

@@ -12,6 +12,7 @@ use bitcoin::{
 
 type BAddress = BAddressGeneric<NetworkChecked>;
 
+// TODO: Add a new so you can't create an address which can't be encoded
 #[derive(Clone, Eq, Debug)]
 pub struct Address(pub BAddress);