From 603cc2a24c01d3461322d8d7993775dbe1c787e8 Mon Sep 17 00:00:00 2001
From: Kartik Agarwala <agarwala.kartik@gmail.com>
Date: Mon, 1 Jul 2024 17:52:23 +0530
Subject: [PATCH] Check lengths after deserialization in cmpctblocks target

---
 fuzzer.cpp                  |  3 +++
 rust_bitcoin_lib/src/lib.rs |  8 ++++----
 targets/cmpctblocks.cpp     | 17 +++++++++--------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/fuzzer.cpp b/fuzzer.cpp
index a82ead0..c116a62 100644
--- a/fuzzer.cpp
+++ b/fuzzer.cpp
@@ -14,6 +14,7 @@
 #include "targets/psbt.h"
 #include "targets/script.h"
 #include "targets/addrv2.h"
+#include "targets/cmpctblocks.h"
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     FuzzedDataProvider provider(data, size);
@@ -37,6 +38,8 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
         Script(provider);
     } else if (target == "addrv2") {
         Addrv2(provider);
+    } else if (target == "cmpctblocks") {
+        CmpctBlocks(provider);
     }
 
     return 0; // Values other than 0 and -1 are reserved for future use.
diff --git a/rust_bitcoin_lib/src/lib.rs b/rust_bitcoin_lib/src/lib.rs
index c249cc9..6f44030 100644
--- a/rust_bitcoin_lib/src/lib.rs
+++ b/rust_bitcoin_lib/src/lib.rs
@@ -137,20 +137,20 @@ pub unsafe extern "C" fn rust_bitcoin_addrv2(data: *const u8, len: usize, actual
 }
 
 #[no_mangle]
-pub unsafe extern "C" fn rust_bitcoin_cmpctblocks(data: *const u8, len: usize) -> *mut c_char {
+pub unsafe extern "C" fn rust_bitcoin_cmpctblocks(data: *const u8, len: usize) -> i32 {
     // Safety: Ensure that the data pointer is valid for the given length
     let data_slice = slice::from_raw_parts(data, len);
 
     let res = deserialize_partial::<HeaderAndShortIds>(data_slice);
 
     match res {
-        Ok(_) => return str_to_c_string("0"),
+        Ok(block) => return (block.0.prefilled_txs.len() + block.0.short_ids.len()).try_into().unwrap(),
         Err(err) => {
             if err.to_string().starts_with("unsupported segwit version") {
-                return str_to_c_string("unsupported segwit version")
+                return -2;
             }
 
-            return str_to_c_string("1")
+            return -1;
         }
     }
 }
diff --git a/targets/cmpctblocks.cpp b/targets/cmpctblocks.cpp
index b8ffded..6e34122 100644
--- a/targets/cmpctblocks.cpp
+++ b/targets/cmpctblocks.cpp
@@ -1,13 +1,14 @@
 #include <fuzzer/FuzzedDataProvider.h>
 #include <string>
+#include <iostream>
 
 #include "cmpctblocks.h"
 #include "bitcoin/src/blockencodings.h"
 #include "bitcoin/src/streams.h"
 
-extern "C" char* rust_bitcoin_cmpctblocks(uint8_t *data, size_t len);
+extern "C" int rust_bitcoin_cmpctblocks(uint8_t *data, size_t len);
 
-std::string CmpctBlocksCore(Span<const uint8_t> buffer) 
+int CmpctBlocksCore(Span<const uint8_t> buffer) 
 {
     DataStream ds{buffer};
     CBlockHeaderAndShortTxIDs block_header_and_short_txids;
@@ -15,19 +16,19 @@ std::string CmpctBlocksCore(Span<const uint8_t> buffer)
         ds >> block_header_and_short_txids;
     } catch (const std::ios_base::failure& e) {
         if (std::string(e.what()).find("Superfluous witness record") != std::string::npos)
-            return "Superfluous witness record";
-        return "1";
+            return -2;
+        return -1;
     }
-    return "0";
+    return block_header_and_short_txids.BlockTxCount();
 }
 
 void CmpctBlocks(FuzzedDataProvider& provider)
 {
     std::vector<uint8_t> buffer{provider.ConsumeRemainingBytes<uint8_t>()};
-    std::string core{CmpctBlocksCore(buffer)};
-    std::string rust_bitcoin{rust_bitcoin_cmpctblocks(buffer.data(), buffer.size())};
+    int core{CmpctBlocksCore(buffer)};
+    int rust_bitcoin{rust_bitcoin_cmpctblocks(buffer.data(), buffer.size())};
 
-    if (core == "Superfluous witness record" || rust_bitcoin == "unsupported segwit version")
+    if (core == -2 || rust_bitcoin == -2)
         return;
 
     assert(core == rust_bitcoin);