From d03ef96126e7f43baff867b0f299553e3bfc1c00 Mon Sep 17 00:00:00 2001
From: Kartik Agarwala <agarwala.kartik@gmail.com>
Date: Mon, 1 Jul 2024 16:00:38 +0530
Subject: [PATCH] Add target for compact block messages

---
 Makefile                            |  2 +-
 rust_bitcoin_lib/rust_bitcoin_lib.h |  1 +
 rust_bitcoin_lib/src/lib.rs         | 20 +++++++++++++++++
 targets/cmpctblocks.cpp             | 34 +++++++++++++++++++++++++++++
 targets/cmpctblocks.h               |  7 ++++++
 5 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 targets/cmpctblocks.cpp
 create mode 100644 targets/cmpctblocks.h

diff --git a/Makefile b/Makefile
index cb7761b..557e266 100644
--- a/Makefile
+++ b/Makefile
@@ -9,7 +9,7 @@ INCPATHS:=  $(foreach dir,$(INCLUDES),-I$(dir))
 LIBPATHS:=  $(foreach lib,$(LIB_DIR),-L$(lib))
 CXXFLAGS:=  -O3 -g0 -Wall -fsanitize=fuzzer -DHAVE_GMTIME_R=1 -std=c++20 -march=native $(INCPATHS)
 ORIGLDFLAGS := $(LDFLAGS) # need to save a copy of ld flags as these get modified below
-LDFLAGS :=  $(LIBPATHS) -lbtcd_wrapper -lrust_bitcoin_lib -lbitcoin_common -lbitcoin_util -lbitcoinkernel -lsecp256k1 -lpthread -ldl
+LDFLAGS :=  $(LIBPATHS) -lbtcd_wrapper -lrust_bitcoin_lib -lbitcoin_node -lbitcoin_common -lbitcoin_util -lbitcoinkernel -lunivalue -lsecp256k1 -lpthread -ldl
 
 ifeq ($(UNAME_S),Darwin)
 LDFLAGS += -framework CoreFoundation -Wl,-ld_classic
diff --git a/rust_bitcoin_lib/rust_bitcoin_lib.h b/rust_bitcoin_lib/rust_bitcoin_lib.h
index 3e478d2..edec187 100644
--- a/rust_bitcoin_lib/rust_bitcoin_lib.h
+++ b/rust_bitcoin_lib/rust_bitcoin_lib.h
@@ -7,3 +7,4 @@ extern char* rust_miniscript_from_str_check_key(const char* miniscript_str);
 extern char* rust_bitcoin_des_block(const uint8_t *data, size_t len);
 extern char* rust_bitcoin_prefilledtransaction(const uint8_t *data, size_t len);
 extern bool rust_bitcoin_addrv2(uint8_t *data, size_t len, uint64_t *count);
+extern char* rust_bitcoin_cmpctblocks(uint8_t *data, size_t len);
\ No newline at end of file
diff --git a/rust_bitcoin_lib/src/lib.rs b/rust_bitcoin_lib/src/lib.rs
index 361c0aa..c249cc9 100644
--- a/rust_bitcoin_lib/src/lib.rs
+++ b/rust_bitcoin_lib/src/lib.rs
@@ -6,6 +6,7 @@ use std::str::FromStr;
 use std::str::Utf8Error;
 
 use bitcoin::bip152::PrefilledTransaction;
+use bitcoin::bip152::HeaderAndShortIds;
 use bitcoin::consensus::deserialize_partial;
 use bitcoin::consensus::encode;
 use bitcoin::Block;
@@ -135,6 +136,25 @@ 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 {
+    // 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"),
+        Err(err) => {
+            if err.to_string().starts_with("unsupported segwit version") {
+                return str_to_c_string("unsupported segwit version")
+            }
+
+            return str_to_c_string("1")
+        }
+    }
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn rust_miniscript_from_str_check_key(input: *const c_char) -> *mut c_char {
     let Ok(desc) = c_str_to_str(input) else {
diff --git a/targets/cmpctblocks.cpp b/targets/cmpctblocks.cpp
new file mode 100644
index 0000000..b8ffded
--- /dev/null
+++ b/targets/cmpctblocks.cpp
@@ -0,0 +1,34 @@
+#include <fuzzer/FuzzedDataProvider.h>
+#include <string>
+
+#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);
+
+std::string CmpctBlocksCore(Span<const uint8_t> buffer) 
+{
+    DataStream ds{buffer};
+    CBlockHeaderAndShortTxIDs block_header_and_short_txids;
+    try {
+        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 "0";
+}
+
+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())};
+
+    if (core == "Superfluous witness record" || rust_bitcoin == "unsupported segwit version")
+        return;
+
+    assert(core == rust_bitcoin);
+}
\ No newline at end of file
diff --git a/targets/cmpctblocks.h b/targets/cmpctblocks.h
new file mode 100644
index 0000000..602cd3e
--- /dev/null
+++ b/targets/cmpctblocks.h
@@ -0,0 +1,7 @@
+#ifndef CMPCT_BLOCKS_H
+#define CMPCT_BLOCKS_H
+
+#include <fuzzer/FuzzedDataProvider.h>
+
+void CmpctBlocks(FuzzedDataProvider& provider);
+#endif
\ No newline at end of file