From 0665c9e854f28d0a171c39b5c0be00d19d3a9062 Mon Sep 17 00:00:00 2001
From: Alekos Filini <alekos.filini@gmail.com>
Date: Fri, 7 Aug 2020 15:35:14 +0200
Subject: [PATCH] [wallet] TxOrdering, shuffle/bip69 support

---
 src/wallet/mod.rs        |   8 +--
 src/wallet/tx_builder.rs | 136 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 133 insertions(+), 11 deletions(-)

diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index 609f2f45..fe1255a5 100644
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -252,12 +252,8 @@ where
             return Err(Error::InsufficientFunds); // TODO: or OutputBelowDustLimit?
         }
 
-        if builder.shuffle_outputs.unwrap_or(true) {
-            use rand::seq::SliceRandom;
-
-            let mut rng = rand::thread_rng();
-            tx.output.shuffle(&mut rng);
-        }
+        // sort input/outputs according to the chosen algorithm
+        builder.ordering.modify_tx(&mut tx);
 
         let txid = tx.txid();
         let mut psbt = PSBT::from_unsigned_tx(tx)?;
diff --git a/src/wallet/tx_builder.rs b/src/wallet/tx_builder.rs
index b644660a..74c69468 100644
--- a/src/wallet/tx_builder.rs
+++ b/src/wallet/tx_builder.rs
@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
+use std::default::Default;
 
-use bitcoin::{Address, OutPoint, SigHashType};
+use bitcoin::{Address, OutPoint, SigHashType, Transaction};
 
 use super::coin_selection::{CoinSelectionAlgorithm, DefaultCoinSelectionAlgorithm};
 use super::utils::FeeRate;
@@ -15,7 +16,8 @@ pub struct TxBuilder<Cs: CoinSelectionAlgorithm> {
     pub(crate) utxos: Option<Vec<OutPoint>>,
     pub(crate) unspendable: Option<Vec<OutPoint>>,
     pub(crate) sighash: Option<SigHashType>,
-    pub(crate) shuffle_outputs: Option<bool>,
+    pub(crate) ordering: TxOrdering,
+    pub(crate) locktime: Option<u32>,
     pub(crate) coin_selection: Cs,
 }
 
@@ -80,8 +82,13 @@ impl<Cs: CoinSelectionAlgorithm> TxBuilder<Cs> {
         self
     }
 
-    pub fn do_not_shuffle_outputs(mut self) -> Self {
-        self.shuffle_outputs = Some(false);
+    pub fn ordering(mut self, ordering: TxOrdering) -> Self {
+        self.ordering = ordering;
+        self
+    }
+
+    pub fn nlocktime(mut self, locktime: u32) -> Self {
+        self.locktime = Some(locktime);
         self
     }
 
@@ -94,8 +101,127 @@ impl<Cs: CoinSelectionAlgorithm> TxBuilder<Cs> {
             utxos: self.utxos,
             unspendable: self.unspendable,
             sighash: self.sighash,
-            shuffle_outputs: self.shuffle_outputs,
+            ordering: self.ordering,
+            locktime: self.locktime,
             coin_selection,
         }
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+pub enum TxOrdering {
+    Shuffle,
+    Untouched,
+    BIP69Lexicographic,
+}
+
+impl Default for TxOrdering {
+    fn default() -> Self {
+        TxOrdering::Shuffle
+    }
+}
+
+impl TxOrdering {
+    pub fn modify_tx(&self, tx: &mut Transaction) {
+        match self {
+            TxOrdering::Untouched => {}
+            TxOrdering::Shuffle => {
+                use rand::seq::SliceRandom;
+                #[cfg(test)]
+                use rand::SeedableRng;
+
+                #[cfg(not(test))]
+                let mut rng = rand::thread_rng();
+                #[cfg(test)]
+                let mut rng = rand::rngs::StdRng::seed_from_u64(0);
+
+                tx.output.shuffle(&mut rng);
+            }
+            TxOrdering::BIP69Lexicographic => {
+                tx.input.sort_unstable_by_key(|txin| {
+                    (txin.previous_output.txid, txin.previous_output.vout)
+                });
+                tx.output
+                    .sort_unstable_by_key(|txout| (txout.value, txout.script_pubkey.clone()));
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    const ORDERING_TEST_TX: &'static str = "0200000003c26f3eb7932f7acddc5ddd26602b77e7516079b03090a16e2c2f54\
+                                            85d1fd600f0100000000ffffffffc26f3eb7932f7acddc5ddd26602b77e75160\
+                                            79b03090a16e2c2f5485d1fd600f0000000000ffffffff571fb3e02278217852\
+                                            dd5d299947e2b7354a639adc32ec1fa7b82cfb5dec530e0500000000ffffffff\
+                                            03e80300000000000002aaeee80300000000000001aa200300000000000001ff\
+                                            00000000";
+    macro_rules! ordering_test_tx {
+        () => {
+            deserialize::<bitcoin::Transaction>(&Vec::<u8>::from_hex(ORDERING_TEST_TX).unwrap())
+                .unwrap()
+        };
+    }
+
+    use bitcoin::consensus::deserialize;
+    use bitcoin::hashes::hex::FromHex;
+
+    use super::*;
+
+    #[test]
+    fn test_output_ordering_untouched() {
+        let original_tx = ordering_test_tx!();
+        let mut tx = original_tx.clone();
+
+        TxOrdering::Untouched.modify_tx(&mut tx);
+
+        assert_eq!(original_tx, tx);
+    }
+
+    #[test]
+    fn test_output_ordering_shuffle() {
+        let original_tx = ordering_test_tx!();
+        let mut tx = original_tx.clone();
+
+        TxOrdering::Shuffle.modify_tx(&mut tx);
+
+        assert_eq!(original_tx.input, tx.input);
+        assert_ne!(original_tx.output, tx.output);
+    }
+
+    #[test]
+    fn test_output_ordering_bip69() {
+        use std::str::FromStr;
+
+        let original_tx = ordering_test_tx!();
+        let mut tx = original_tx.clone();
+
+        TxOrdering::BIP69Lexicographic.modify_tx(&mut tx);
+
+        assert_eq!(
+            tx.input[0].previous_output,
+            bitcoin::OutPoint::from_str(
+                "0e53ec5dfb2cb8a71fec32dc9a634a35b7e24799295ddd5278217822e0b31f57:5"
+            )
+            .unwrap()
+        );
+        assert_eq!(
+            tx.input[1].previous_output,
+            bitcoin::OutPoint::from_str(
+                "0f60fdd185542f2c6ea19030b0796051e7772b6026dd5ddccd7a2f93b73e6fc2:0"
+            )
+            .unwrap()
+        );
+        assert_eq!(
+            tx.input[2].previous_output,
+            bitcoin::OutPoint::from_str(
+                "0f60fdd185542f2c6ea19030b0796051e7772b6026dd5ddccd7a2f93b73e6fc2:1"
+            )
+            .unwrap()
+        );
+
+        assert_eq!(tx.output[0].value, 800);
+        assert_eq!(tx.output[1].script_pubkey, From::from(vec![0xAA]));
+        assert_eq!(tx.output[2].script_pubkey, From::from(vec![0xAA, 0xEE]));
+    }
+}