Allow restoring wallet with custom timestamp (ordinals#4065)

src/subcommand/wallet/restore.rs

@@ -1,11 +1,32 @@
 use super::*;
 
+#[derive(Debug, Clone)]
+pub(crate) struct Timestamp(bitcoincore_rpc::json::Timestamp);
+
+impl FromStr for Timestamp {
+  type Err = Error;
+
+  fn from_str(s: &str) -> Result<Self, Self::Err> {
+    Ok(if s == "now" {
+      Self(bitcoincore_rpc::json::Timestamp::Now)
+    } else {
+      Self(bitcoincore_rpc::json::Timestamp::Time(s.parse()?))
+    })
+  }
+}
+
 #[derive(Debug, Parser)]
 pub(crate) struct Restore {
   #[clap(value_enum, long, help = "Restore wallet from <SOURCE> on stdin.")]
   from: Source,
-  #[arg(long, help = "Use <PASSPHRASE> when deriving wallet")]
+  #[arg(long, help = "Use <PASSPHRASE> when deriving wallet.")]
   pub(crate) passphrase: Option<String>,
+  #[arg(
+    long,
+    help = "Scan chain from <TIMESTAMP> onwards. Can be a unix timestamp in \
+    seconds or the string `now`, to skip scanning"
+  )]
+  pub(crate) timestamp: Option<Timestamp>,
 }
 
 #[derive(clap::ValueEnum, Debug, Clone)]
@@ -31,10 +52,17 @@ impl Restore {
     match self.from {
       Source::Descriptor => {
         io::stdin().read_to_string(&mut buffer)?;
+
         ensure!(
           self.passphrase.is_none(),
           "descriptor does not take a passphrase"
         );
+
+        ensure!(
+          self.timestamp.is_none(),
+          "descriptor does not take a timestamp"
+        );
+
         let wallet_descriptors: ListDescriptorsResult = serde_json::from_str(&buffer)?;
         Wallet::initialize_from_descriptors(name, settings, wallet_descriptors.descriptors)?;
       }
@@ -45,7 +73,10 @@ impl Restore {
           name,
           settings,
           mnemonic.to_seed(self.passphrase.unwrap_or_default()),
-          bitcoincore_rpc::json::Timestamp::Time(0),
+          self
+            .timestamp
+            .unwrap_or(Timestamp(bitcoincore_rpc::json::Timestamp::Time(0)))
+            .0,
         )?;
       }
     }

src/wallet.rs

@@ -2,10 +2,10 @@ use {
   super::*,
   base64::{self, Engine},
   batch::ParentInfo,
-  bitcoin::secp256k1::{All, Secp256k1},
   bitcoin::{
-    bip32::{ChildNumber, DerivationPath, Fingerprint, Xpriv},
+    bip32::{ChildNumber, DerivationPath, Xpriv},
     psbt::Psbt,
+    secp256k1::Secp256k1,
   },
   bitcoincore_rpc::json::ImportDescriptors,
   entry::{EtchingEntry, EtchingEntryValue},
@@ -526,57 +526,38 @@ impl Wallet {
 
     let derived_private_key = master_private_key.derive_priv(&secp, &derivation_path)?;
 
+    let mut descriptors = Vec::new();
     for change in [false, true] {
-      Self::derive_and_import_descriptor(
-        name.clone(),
-        settings,
-        &secp,
-        (fingerprint, derivation_path.clone()),
-        derived_private_key,
-        change,
-        timestamp,
-      )?;
-    }
-
-    Ok(())
-  }
-
-  fn derive_and_import_descriptor(
-    name: String,
-    settings: &Settings,
-    secp: &Secp256k1<All>,
-    origin: (Fingerprint, DerivationPath),
-    derived_private_key: Xpriv,
-    change: bool,
-    timestamp: bitcoincore_rpc::json::Timestamp,
-  ) -> Result {
-    let secret_key = DescriptorSecretKey::XPrv(DescriptorXKey {
-      origin: Some(origin),
-      xkey: derived_private_key,
-      derivation_path: DerivationPath::master().child(ChildNumber::Normal {
-        index: change.into(),
-      }),
-      wildcard: Wildcard::Unhardened,
-    });
+      let secret_key = DescriptorSecretKey::XPrv(DescriptorXKey {
+        origin: Some((fingerprint, derivation_path.clone())),
+        xkey: derived_private_key,
+        derivation_path: DerivationPath::master().child(ChildNumber::Normal {
+          index: change.into(),
+        }),
+        wildcard: Wildcard::Unhardened,
+      });
 
-    let public_key = secret_key.to_public(secp)?;
+      let public_key = secret_key.to_public(&secp)?;
 
-    let mut key_map = BTreeMap::new();
-    key_map.insert(public_key.clone(), secret_key);
+      let mut key_map = BTreeMap::new();
+      key_map.insert(public_key.clone(), secret_key);
 
-    let descriptor = miniscript::descriptor::Descriptor::new_tr(public_key, None)?;
+      let descriptor = miniscript::descriptor::Descriptor::new_tr(public_key, None)?;
 
-    settings
-      .bitcoin_rpc_client(Some(name.clone()))?
-      .import_descriptors(ImportDescriptors {
+      descriptors.push(ImportDescriptors {
         descriptor: descriptor.to_string_with_secret(&key_map),
         timestamp,
         active: Some(true),
         range: None,
         next_index: None,
         internal: Some(change),
         label: None,
-      })?;
+      });
+    }
+
+    settings
+      .bitcoin_rpc_client(Some(name.clone()))?
+      .call::<serde_json::Value>("importdescriptors", &[serde_json::to_value(descriptors)?])?;
 
     Ok(())
   }

tests/wallet/restore.rs

@@ -222,3 +222,149 @@ fn passphrase_conflicts_with_descriptor() {
   .expected_stderr("error: descriptor does not take a passphrase\n")
   .run_and_extract_stdout();
 }
+
+#[test]
+fn timestamp_conflicts_with_descriptor() {
+  let core = mockcore::spawn();
+  let ord = TestServer::spawn(&core);
+
+  CommandBuilder::new([
+    "wallet",
+    "restore",
+    "--from",
+    "descriptor",
+    "--timestamp",
+    "now",
+  ])
+  .stdin("".into())
+  .core(&core)
+  .ord(&ord)
+  .expected_exit_code(1)
+  .expected_stderr("error: descriptor does not take a timestamp\n")
+  .run_and_extract_stdout();
+}
+
+#[test]
+fn restore_with_now_timestamp() {
+  let mnemonic = {
+    let core = mockcore::spawn();
+
+    let create::Output { mnemonic, .. } = CommandBuilder::new(["wallet", "create"])
+      .core(&core)
+      .run_and_deserialize_output();
+
+    mnemonic
+  };
+
+  let core = mockcore::spawn();
+  let ord = TestServer::spawn(&core);
+
+  CommandBuilder::new([
+    "wallet",
+    "restore",
+    "--from",
+    "mnemonic",
+    "--timestamp",
+    "now",
+  ])
+  .stdin(mnemonic.to_string().into())
+  .core(&core)
+  .run_and_extract_stdout();
+
+  let output = CommandBuilder::new("wallet dump")
+    .core(&core)
+    .ord(&ord)
+    .stderr_regex(".*")
+    .run_and_deserialize_output::<ListDescriptorsResult>();
+
+  assert!(output
+    .descriptors
+    .iter()
+    .all(|descriptor| match descriptor.timestamp {
+      bitcoincore_rpc::json::Timestamp::Now => true,
+      bitcoincore_rpc::json::Timestamp::Time(time) =>
+        time.abs_diff(
+          std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_secs()
+        ) <= 5,
+    }));
+}
+
+#[test]
+fn restore_with_no_timestamp_defaults_to_0() {
+  let mnemonic = {
+    let core = mockcore::spawn();
+
+    let create::Output { mnemonic, .. } = CommandBuilder::new(["wallet", "create"])
+      .core(&core)
+      .run_and_deserialize_output();
+
+    mnemonic
+  };
+
+  let core = mockcore::spawn();
+  let ord = TestServer::spawn(&core);
+
+  CommandBuilder::new(["wallet", "restore", "--from", "mnemonic"])
+    .stdin(mnemonic.to_string().into())
+    .core(&core)
+    .run_and_extract_stdout();
+
+  let output = CommandBuilder::new("wallet dump")
+    .core(&core)
+    .ord(&ord)
+    .stderr_regex(".*")
+    .run_and_deserialize_output::<ListDescriptorsResult>();
+
+  assert!(output
+    .descriptors
+    .iter()
+    .all(|descriptor| match descriptor.timestamp {
+      bitcoincore_rpc::json::Timestamp::Now => false,
+      bitcoincore_rpc::json::Timestamp::Time(time) => time == 0,
+    }));
+}
+
+#[test]
+fn restore_with_timestamp() {
+  let mnemonic = {
+    let core = mockcore::spawn();
+
+    let create::Output { mnemonic, .. } = CommandBuilder::new(["wallet", "create"])
+      .core(&core)
+      .run_and_deserialize_output();
+
+    mnemonic
+  };
+
+  let core = mockcore::spawn();
+  let ord = TestServer::spawn(&core);
+
+  CommandBuilder::new([
+    "wallet",
+    "restore",
+    "--from",
+    "mnemonic",
+    "--timestamp",
+    "123456789",
+  ])
+  .stdin(mnemonic.to_string().into())
+  .core(&core)
+  .run_and_extract_stdout();
+
+  let output = CommandBuilder::new("wallet dump")
+    .core(&core)
+    .ord(&ord)
+    .stderr_regex(".*")
+    .run_and_deserialize_output::<ListDescriptorsResult>();
+
+  assert!(output
+    .descriptors
+    .iter()
+    .all(|descriptor| match descriptor.timestamp {
+      bitcoincore_rpc::json::Timestamp::Now => false,
+      bitcoincore_rpc::json::Timestamp::Time(time) => time == 123456789,
+    }));
+}