From a7d478024ed6eba7a8564bce33a8afa757521d6f Mon Sep 17 00:00:00 2001
From: liam <31192478+terror@users.noreply.github.com>
Date: Thu, 4 Aug 2022 18:42:13 -0400
Subject: [PATCH] Add network option (#274)

---
 src/arguments.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
 src/index.rs     | 16 +++----------
 src/main.rs      |  1 +
 src/options.rs   | 40 +++++++++++++++++++++++++++++++
 tests/lib.rs     | 11 +++++++++
 5 files changed, 116 insertions(+), 13 deletions(-)

diff --git a/src/arguments.rs b/src/arguments.rs
index 6be3eeb4ec..7a6b9c1129 100644
--- a/src/arguments.rs
+++ b/src/arguments.rs
@@ -14,3 +14,64 @@ impl Arguments {
     self.subcommand.run(self.options)
   }
 }
+
+#[cfg(test)]
+mod tests {
+  use {super::*, std::path::Path};
+
+  #[test]
+  fn rpc_url_overrides_network() {
+    assert_eq!(
+      Arguments::try_parse_from(&[
+        "ord",
+        "--rpc-url=127.0.0.1:1234",
+        "--network=signet",
+        "index"
+      ])
+      .unwrap()
+      .options
+      .rpc_url(),
+      "127.0.0.1:1234"
+    );
+  }
+
+  #[test]
+  fn cookie_file_overrides_network() {
+    assert_eq!(
+      Arguments::try_parse_from(&["ord", "--cookie-file=/foo/bar", "--network=signet", "index"])
+        .unwrap()
+        .options
+        .cookie_file()
+        .unwrap(),
+      Path::new("/foo/bar")
+    );
+  }
+
+  #[test]
+  fn use_default_network() {
+    let arguments = Arguments::try_parse_from(&["ord", "index"]).unwrap();
+
+    assert_eq!(arguments.options.rpc_url(), "127.0.0.1:8333");
+
+    assert!(arguments
+      .options
+      .cookie_file()
+      .unwrap()
+      .ends_with(".cookie"));
+  }
+
+  #[test]
+  fn uses_network_defaults() {
+    let arguments = Arguments::try_parse_from(&["ord", "--network=signet", "index"]).unwrap();
+
+    assert_eq!(arguments.options.rpc_url(), "127.0.0.1:38333");
+
+    assert!(arguments
+      .options
+      .cookie_file()
+      .unwrap()
+      .display()
+      .to_string()
+      .ends_with("/signet/.cookie"))
+  }
+}
diff --git a/src/index.rs b/src/index.rs
index c6ecadf43b..cc02c604e3 100644
--- a/src/index.rs
+++ b/src/index.rs
@@ -1,6 +1,6 @@
 use {
   super::*,
-  bitcoincore_rpc::{Auth, Client, RpcApi},
+  bitcoincore_rpc::{Client, RpcApi},
   rayon::iter::{IntoParallelRefIterator, ParallelIterator},
 };
 
@@ -15,18 +15,8 @@ pub(crate) struct Index {
 
 impl Index {
   pub(crate) fn open(options: &Options) -> Result<Self> {
-    let client = Client::new(
-      options
-        .rpc_url
-        .as_ref()
-        .ok_or_else(|| anyhow!("This command requires `--rpc-url`"))?,
-      options
-        .cookie_file
-        .as_ref()
-        .map(|path| Auth::CookieFile(path.clone()))
-        .unwrap_or(Auth::None),
-    )
-    .context("Failed to connect to RPC URL")?;
+    let client = Client::new(&options.rpc_url(), Auth::CookieFile(options.cookie_file()?))
+      .context("Failed to connect to RPC URL")?;
 
     let database = match unsafe { redb::Database::open("index.redb") } {
       Ok(database) => database,
diff --git a/src/main.rs b/src/main.rs
index 2920dced3b..2b12b022a4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -29,6 +29,7 @@ use {
     util::key::PrivateKey,
     Block, Network, OutPoint, Transaction, Txid,
   },
+  bitcoincore_rpc::Auth,
   chrono::{DateTime, NaiveDateTime, Utc},
   clap::Parser,
   derive_more::{Display, FromStr},
diff --git a/src/options.rs b/src/options.rs
index 99b1e75028..1e561053d9 100644
--- a/src/options.rs
+++ b/src/options.rs
@@ -8,4 +8,44 @@ pub(crate) struct Options {
   pub(crate) cookie_file: Option<PathBuf>,
   #[clap(long)]
   pub(crate) rpc_url: Option<String>,
+  #[clap(long, default_value = "bitcoin")]
+  pub(crate) network: Network,
+}
+
+impl Options {
+  pub(crate) fn rpc_url(&self) -> String {
+    self
+      .rpc_url
+      .as_ref()
+      .unwrap_or(&format!(
+        "127.0.0.1:{}",
+        match self.network {
+          Network::Bitcoin => "8333",
+          Network::Regtest => "18443",
+          Network::Signet => "38333",
+          Network::Testnet => "18332",
+        }
+      ))
+      .into()
+  }
+
+  pub(crate) fn cookie_file(&self) -> Result<PathBuf> {
+    if let Some(cookie_file) = &self.cookie_file {
+      return Ok(cookie_file.clone());
+    }
+
+    let mut path = if cfg!(linux) {
+      dirs::home_dir()
+        .ok_or_else(|| anyhow!("Failed to retrieve home dir"))?
+        .join(".bitcoin")
+    } else {
+      dirs::data_dir().ok_or_else(|| anyhow!("Failed to retrieve data dir"))?
+    };
+
+    if !matches!(self.network, Network::Bitcoin) {
+      path.push(self.network.to_string())
+    }
+
+    Ok(path.join(".cookie"))
+  }
 }
diff --git a/tests/lib.rs b/tests/lib.rs
index 15f0fa4d37..cb444aa05b 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -119,6 +119,16 @@ impl Test {
   }
 
   fn with_tempdir(tempdir: TempDir) -> Self {
+    let cookie_file = tempdir.path().join(".cookie");
+
+    if !cookie_file.exists() {
+      fs::write(
+        cookie_file,
+        "__cookie__:f5c6aedf2ed57e81856202def76bec8cb63f56e06f5cb04eb996eb831248d95d",
+      )
+      .unwrap();
+    }
+
     Self {
       args: Vec::new(),
       envs: Vec::new(),
@@ -254,6 +264,7 @@ impl Test {
       })
       .current_dir(&self.tempdir)
       .arg(format!("--rpc-url=http://127.0.0.1:{rpc_server_port}"))
+      .arg("--cookie-file=.cookie")
       .args(self.args)
       .spawn()?;