From b9baa4691e913536b5722fd0acacb2e0dafffeb1 Mon Sep 17 00:00:00 2001
From: Jonathan Becker <jonathan@jbecker.dev>
Date: Sun, 3 Sep 2023 12:49:44 -0400
Subject: [PATCH] :sparkles: feat: `--truncate-calldata` flag (#145)

---
 heimdall/src/decode/mod.rs    | 20 ++++++++++++++++++--
 heimdall/tests/test_decode.rs |  6 ++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/heimdall/src/decode/mod.rs b/heimdall/src/decode/mod.rs
index f36e44b8..42175508 100644
--- a/heimdall/src/decode/mod.rs
+++ b/heimdall/src/decode/mod.rs
@@ -55,6 +55,10 @@ pub struct DecodeArgs {
     /// When prompted, always select the default value.
     #[clap(long, short)]
     pub default: bool,
+
+    /// Whether to truncate nonstandard sized calldata.
+    #[clap(long, short)]
+    pub truncate_calldata: bool,
 }
 
 #[allow(deprecated)]
@@ -90,7 +94,7 @@ pub fn decode(args: DecodeArgs) {
     }
 
     // normalize
-    let calldata = calldata.replacen("0x", "", 1);
+    let mut calldata = calldata.replacen("0x", "", 1);
 
     // check if the calldata length is a standard length
     if calldata.len() % 2 != 0 {
@@ -99,8 +103,20 @@ pub fn decode(args: DecodeArgs) {
     }
 
     // if calldata isn't a multiple of 64, it may be harder to decode.
-    if calldata[8..].len() % 64 != 0 {
+    if (calldata[8..].len() % 64 != 0) && !args.truncate_calldata {
         logger.warn("calldata is not a standard size. decoding may fail since each word is not exactly 32 bytes long.");
+        logger.warn("if decoding fails, try using the --truncate-calldata flag to truncate the calldata to a standard size.");
+    } else if args.truncate_calldata {
+        logger.warn("calldata is not a standard size. truncating the calldata to a standard size.");
+
+        // get the selector
+        let selector = calldata[0..8].to_owned();
+
+        // truncate calldata to a standard size, removing the selector
+        calldata = calldata[8..][..calldata[8..].len() - (calldata[8..].len() % 64)].to_owned();
+
+        // add the selector back
+        calldata = selector + &calldata;
     }
 
     // parse the two parts of calldata, inputs and selector
diff --git a/heimdall/tests/test_decode.rs b/heimdall/tests/test_decode.rs
index 6f3eb8f1..c224281b 100644
--- a/heimdall/tests/test_decode.rs
+++ b/heimdall/tests/test_decode.rs
@@ -15,6 +15,7 @@ mod benchmark {
                 openai_api_key: String::from(""),
                 explain: false,
                 default: true,
+                truncate_calldata: false,
             };
             heimdall::decode::decode(args)
         }
@@ -32,6 +33,7 @@ mod benchmark {
                 openai_api_key: String::from(""),
                 explain: false,
                 default: true,
+                truncate_calldata: false,
             };
             heimdall::decode::decode(args)
         }
@@ -49,6 +51,7 @@ mod benchmark {
                 openai_api_key: String::from(""),
                 explain: false,
                 default: true,
+                truncate_calldata: false,
             };
             heimdall::decode::decode(args)
         }
@@ -66,6 +69,7 @@ mod benchmark {
                 openai_api_key: String::from(""),
                 explain: false,
                 default: true,
+                truncate_calldata: false,
             };
             heimdall::decode::decode(args)
         }
@@ -82,6 +86,7 @@ mod benchmark {
             openai_api_key: String::from(""),
             explain: false,
             default: true,
+            truncate_calldata: false,
         };
         heimdall::decode::decode(args);
         assert!(true)
@@ -96,6 +101,7 @@ mod benchmark {
             openai_api_key: String::from(""),
             explain: false,
             default: true,
+            truncate_calldata: false,
         };
         heimdall::decode::decode(args);
         assert!(true)