feat: Add decode-key command to http_kv_tool

crates/iota-storage/src/bin/http_kv_tool.rs

@@ -15,118 +15,174 @@ use iota_types::{
     messages_checkpoint::CheckpointSequenceNumber,
 };
 
-// Command line options are:
-// --base-url <url> - the base URL of the HTTP server
-// --digest <digest> - the digest of the key being fetched
-// --type <fx|tx|ev> - the type of key being fetched
 #[derive(Parser)]
-struct Options {
-    #[arg(short, long)]
-    base_url: String,
-
-    #[arg(short, long)]
-    digest: Vec<String>,
-
-    #[arg(short, long)]
-    seq: Vec<String>,
-
-    // must be either 'tx', 'fx','ob','events', or 'ckpt_contents'
-    // default value of 'tx'
-    #[arg(short, long, default_value = "tx")]
-    type_: String,
+enum Command {
+    Fetch {
+        #[arg(short, long)]
+        base_url: String,
+
+        #[arg(short, long)]
+        digest: Vec<String>,
+
+        #[arg(short, long)]
+        seq: Vec<String>,
+
+        // must be either 'tx', 'fx','ob','events','cs', or 'cc'
+        // default value of 'tx'
+        #[arg(short, long, default_value = "tx")]
+        type_: String,
+    },
+
+    DecodeKey {
+        #[arg(short, long)]
+        url: String,
+    },
 }
 
-#[tokio::main]
-async fn main() {
-    let _guard = telemetry_subscribers::TelemetryConfig::new()
-        .with_env()
-        .init();
-
-    let options = Options::parse();
-
-    let metrics = KeyValueStoreMetrics::new_for_tests();
-    let http_kv = Arc::new(HttpKVStore::new(&options.base_url, 100, metrics).unwrap());
-    let kv =
-        TransactionKeyValueStore::new("http_kv", KeyValueStoreMetrics::new_for_tests(), http_kv);
-
-    let seqs: Vec<_> = options
-        .seq
-        .into_iter()
-        .map(|s| {
-            CheckpointSequenceNumber::from_str(&s).expect("invalid checkpoint sequence number")
-        })
-        .collect();
-
-    // verify that type is valid
-    match options.type_.as_str() {
-        "tx" | "fx" => {
-            let digests: Vec<_> = options
-                .digest
-                .into_iter()
-                .map(|digest| {
-                    TransactionDigest::from_str(&digest).expect("invalid transaction digest")
-                })
-                .collect();
-
-            if options.type_ == "tx" {
-                let tx = kv.multi_get_tx(&digests).await.unwrap();
-                for (digest, tx) in digests.iter().zip(tx.iter()) {
-                    println!("fetched tx: {:?} {:?}", digest, tx);
-                }
-            } else {
-                let fx = kv.multi_get_fx_by_tx_digest(&digests).await.unwrap();
-                for (digest, fx) in digests.iter().zip(fx.iter()) {
-                    println!("fetched fx: {:?} {:?}", digest, fx);
+impl Command {
+    async fn execute(self) -> anyhow::Result<()> {
+        match self {
+            Command::Fetch {
+                base_url,
+                digest,
+                seq,
+                type_,
+            } => {
+                let metrics = KeyValueStoreMetrics::new_for_tests();
+                let http_kv = Arc::new(HttpKVStore::new(&base_url, 100, metrics).unwrap());
+                let kv = TransactionKeyValueStore::new(
+                    "http_kv",
+                    KeyValueStoreMetrics::new_for_tests(),
+                    http_kv,
+                );
+
+                let seqs: Vec<_> = seq
+                    .into_iter()
+                    .map(|s| {
+                        CheckpointSequenceNumber::from_str(&s)
+                            .expect("invalid checkpoint sequence number")
+                    })
+                    .collect();
+
+                let item_type = ItemType::from_str(&type_).expect("invalid item type");
+
+                match item_type {
+                    ItemType::Transaction | ItemType::TransactionEffects => {
+                        let digests: Vec<_> = digest
+                            .into_iter()
+                            .map(|digest| {
+                                TransactionDigest::from_str(&digest)
+                                    .expect("invalid transaction digest")
+                            })
+                            .collect();
+
+                        if item_type == ItemType::Transaction {
+                            let tx = kv.multi_get_tx(&digests).await.unwrap();
+                            for (digest, tx) in digests.iter().zip(tx.iter()) {
+                                println!("fetched tx: {:?} {:?}", digest, tx);
+                            }
+                        } else {
+                            let fx = kv.multi_get_fx_by_tx_digest(&digests).await.unwrap();
+                            for (digest, fx) in digests.iter().zip(fx.iter()) {
+                                println!("fetched fx: {:?} {:?}", digest, fx);
+                            }
+                        }
+                    }
+
+                    ItemType::CheckpointContents => {
+                        let ckpts = kv.multi_get_checkpoints(&[], &seqs, &[]).await.unwrap();
+
+                        for (seq, ckpt) in seqs.iter().zip(ckpts.1.iter()) {
+                            // populate digest before printing
+                            ckpt.as_ref().map(|c| c.digest());
+                            println!("fetched ckpt contents: {:?} {:?}", seq, ckpt);
+                        }
+                    }
+
+                    ItemType::CheckpointSummary => {
+                        let digests: Vec<_> = digest
+                            .into_iter()
+                            .map(|s| {
+                                CheckpointDigest::from_str(&s).expect("invalid checkpoint digest")
+                            })
+                            .collect();
+
+                        let ckpts = kv
+                            .multi_get_checkpoints(&seqs, &[], &digests)
+                            .await
+                            .unwrap();
+
+                        for (seq, ckpt) in seqs.iter().zip(ckpts.0.iter()) {
+                            // populate digest before printing
+                            ckpt.as_ref().map(|c| c.digest());
+                            println!("fetched ckpt summary: {:?} {:?}", seq, ckpt);
+                        }
+                        for (digest, ckpt) in digests.iter().zip(ckpts.2.iter()) {
+                            // populate digest before printing
+                            ckpt.as_ref().map(|c| c.digest());
+                            println!("fetched ckpt summary: {:?} {:?}", digest, ckpt);
+                        }
+                    }
+
+                    ItemType::Object => {
+                        let object_id = ObjectID::from_str(&digest[0]).expect("invalid object id");
+                        let object = kv.get_object(object_id, seqs[0].into()).await.unwrap();
+                        println!("fetched object {:?}", object);
+                    }
+
+                    _ => {
+                        println!(
+                            "Invalid key type: {}. Must be one of 'tx', 'fx', or 'ev'.",
+                            type_
+                        );
+                        std::process::exit(1);
+                    }
                 }
+                Ok(())
             }
-        }
+            Command::DecodeKey { url } => {
+                // url may look like
+                // https://{TRANSACTION_KV_STORE_BASE_URL}/jlkqmZbVuunngIyy2vjBOJSETrM56EH_kIc5wuLvDydN_x0GAAAAAA/ob
 
-        "ckpt_contents" => {
-            let ckpts = kv.multi_get_checkpoints(&[], &seqs, &[]).await.unwrap();
+                // extract the digest and type
+                let parts: Vec<_> = url.split('/').collect();
 
-            for (seq, ckpt) in seqs.iter().zip(ckpts.1.iter()) {
-                // populate digest before printing
-                ckpt.as_ref().map(|c| c.digest());
-                println!("fetched ckpt contents: {:?} {:?}", seq, ckpt);
-            }
-        }
+                // its allowed to supply either the whole URL, or the last two pieces
+                if parts.len() < 2 {
+                    println!("Invalid URL: {}", url);
+                    std::process::exit(1);
+                }
 
-        "ckpt_summary" => {
-            let digests: Vec<_> = options
-                .digest
-                .into_iter()
-                .map(|s| CheckpointDigest::from_str(&s).expect("invalid checkpoint digest"))
-                .collect();
-
-            let ckpts = kv
-                .multi_get_checkpoints(&seqs, &[], &digests)
-                .await
-                .unwrap();
-
-            for (seq, ckpt) in seqs.iter().zip(ckpts.0.iter()) {
-                // populate digest before printing
-                ckpt.as_ref().map(|c| c.digest());
-                println!("fetched ckpt summary: {:?} {:?}", seq, ckpt);
-            }
-            for (digest, ckpt) in digests.iter().zip(ckpts.2.iter()) {
-                // populate digest before printing
-                ckpt.as_ref().map(|c| c.digest());
-                println!("fetched ckpt summary: {:?} {:?}", digest, ckpt);
+                let identifier = parts[parts.len() - 2];
+                let type_ = ItemType::from_str(parts[parts.len() - 1]).expect("invalid item type");
+
+                let key = path_elements_to_key(type_, identifier)?;
+                println!("decoded key: {:?}", key);
+                Ok(())
             }
         }
+    }
+}
 
-        "ob" => {
-            let object_id = ObjectID::from_str(&options.digest[0]).expect("invalid object id");
-            let object = kv.get_object(object_id, seqs[0].into()).await.unwrap();
-            println!("fetched object {:?}", object);
-        }
+#[derive(Parser)]
+#[command(
+    name = "http_kv_tool",
+    about = "Utilities for the HTTP key-value store",
+    rename_all = "kebab-case",
+    author,
+    version = "1"
+)]
+struct App {
+    #[command(subcommand)]
+    command: Command,
+}
 
-        _ => {
-            println!(
-                "Invalid key type: {}. Must be one of 'tx', 'fx', or 'ev'.",
-                options.type_
-            );
-            std::process::exit(1);
-        }
-    }
+#[tokio::main]
+async fn main() {
+    let _guard = telemetry_subscribers::TelemetryConfig::new()
+        .with_env()
+        .init();
+
+    let app = App::parse();
+    app.command.execute().await.unwrap();
 }

crates/iota-storage/src/http_key_value_store.rs

@@ -282,6 +282,42 @@ enum Value {
     TxToCheckpoint(CheckpointSequenceNumber),
 }
 
+pub fn path_elements_to_key(type_: ItemType, digest: &str) -> anyhow::Result<Key> {
+    let decoded_digest = base64_url::decode(digest)?;
+
+    match type_ {
+        ItemType::Transaction => Ok(Key::Transaction(TransactionDigest::try_from(
+            decoded_digest,
+        )?)),
+        ItemType::TransactionEffects => Ok(Key::TransactionEffects(TransactionDigest::try_from(
+            decoded_digest,
+        )?)),
+        ItemType::CheckpointContents => {
+            let tagged_key = bcs::from_bytes(&decoded_digest)?;
+            match tagged_key {
+                TaggedKey::CheckpointSequenceNumber(seq) => Ok(Key::CheckpointContents(seq)),
+            }
+        }
+        ItemType::CheckpointSummary => match CheckpointDigest::try_from(decoded_digest.clone()) {
+            Err(_) => {
+                let tagged_key = bcs::from_bytes(&decoded_digest)?;
+                match tagged_key {
+                    TaggedKey::CheckpointSequenceNumber(seq) => Ok(Key::CheckpointSummary(seq)),
+                }
+            }
+            Ok(cs_digest) => Ok(Key::CheckpointSummaryByDigest(cs_digest)),
+        },
+        ItemType::TransactionToCheckpoint => Ok(Key::TransactionToCheckpoint(
+            TransactionDigest::try_from(decoded_digest)?,
+        )),
+        ItemType::Object => {
+            let object_key: ObjectKey = bcs::from_bytes(&decoded_digest)?;
+            Ok(Key::ObjectKey(object_key.0, object_key.1))
+        }
+        _ => Err(anyhow::anyhow!("Invalid ItemType: {:?}", type_)),
+    }
+}
+
 impl HttpKVStore {
     pub fn new_kv(
         base_url: &str,