From fbb312904bba52b5f2606b1d8549211e32be88c2 Mon Sep 17 00:00:00 2001
From: Bruno Garcia <brunoely.gc@gmail.com>
Date: Tue, 18 Jun 2024 10:15:05 -0300
Subject: [PATCH] add target for psbt

---
 Makefile                            |  2 +-
 fuzzer.cpp                          |  3 +++
 rust_bitcoin_lib/rust_bitcoin_lib.h |  1 +
 rust_bitcoin_lib/src/lib.rs         | 20 +++++++++++++++++
 targets/psbt.cpp                    | 33 +++++++++++++++++++++++++++++
 targets/psbt.h                      |  7 ++++++
 6 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 targets/psbt.cpp
 create mode 100644 targets/psbt.h

diff --git a/Makefile b/Makefile
index 730552c..51cc6de 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,6 @@
 CXX      =  clang++
 CC       =  clang
-SOURCES :=  targets/bech32.cpp targets/tx_des.cpp targets/miniscript_string.cpp targets/block_des.cpp targets/prefilledtransaction.cpp
+SOURCES :=  targets/psbt.cpp targets/bech32.cpp targets/tx_des.cpp targets/miniscript_string.cpp targets/block_des.cpp targets/prefilledtransaction.cpp
 INCLUDES =  dependencies/ dependencies/bitcoin/src/ dependencies/bitcoin/src/secp256k1/include
 LIB_DIR  =  dependencies/bitcoin/src/ dependencies/bitcoin/src/.libs dependencies/bitcoin/src/secp256k1/.libs rust_bitcoin_lib/target/debug btcd_lib
 OBJS    :=  $(patsubst %.cpp, build/%.o, $(SOURCES))
diff --git a/fuzzer.cpp b/fuzzer.cpp
index ff93cec..aa3725f 100644
--- a/fuzzer.cpp
+++ b/fuzzer.cpp
@@ -11,6 +11,7 @@
 #include "targets/prefilledtransaction.h"
 #include "targets/tx_des.h"
 #include "targets/bech32.h"
+#include "targets/psbt.h"
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     FuzzedDataProvider provider(data, size);
@@ -28,6 +29,8 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
         TransactionDes(provider);
     } else if (target == "bech32") {
         Bech32(provider);
+    } else if (target == "psbt") {
+        Psbt(provider);
     }
 
     return 0; // Values other than 0 and -1 are reserved for future use.
diff --git a/rust_bitcoin_lib/rust_bitcoin_lib.h b/rust_bitcoin_lib/rust_bitcoin_lib.h
index 2439344..8e0f382 100644
--- a/rust_bitcoin_lib/rust_bitcoin_lib.h
+++ b/rust_bitcoin_lib/rust_bitcoin_lib.h
@@ -1,6 +1,7 @@
 #include <string>
 
 extern bool rust_miniscript_from_str(const char* miniscript_str);
+extern char* rust_bitcoin_psbt(const char* miniscript_str);
 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);
diff --git a/rust_bitcoin_lib/src/lib.rs b/rust_bitcoin_lib/src/lib.rs
index 10bc6e7..269aa8a 100644
--- a/rust_bitcoin_lib/src/lib.rs
+++ b/rust_bitcoin_lib/src/lib.rs
@@ -81,6 +81,26 @@ pub unsafe extern "C" fn rust_miniscript_from_str(input: *const c_char) -> bool
     }
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn rust_bitcoin_psbt(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 psbt: Result<bitcoin::psbt::Psbt, _> = bitcoin::psbt::Psbt::deserialize(data_slice);
+    match psbt {
+        Err(err) => {
+            println!("{}", err.to_string());
+            // Core doesn't check keys and rust-miniscript doesn't support all segwit flags
+            if err.to_string().starts_with("invalid xonly public key") ||
+               err.to_string().starts_with("bitcoin consensus encoding error") {
+                return str_to_c_string(&err.to_string());
+            }
+            str_to_c_string("")
+        },
+        Ok(_) => 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/psbt.cpp b/targets/psbt.cpp
new file mode 100644
index 0000000..72eba33
--- /dev/null
+++ b/targets/psbt.cpp
@@ -0,0 +1,33 @@
+#include <fuzzer/FuzzedDataProvider.h>
+#include <string>
+#include <iostream>
+#include <stdio.h>
+
+#include "psbt.h"
+#include "bitcoin/src/psbt.h"
+#include "bitcoin/src/span.h"
+#include "bitcoin/src/node/psbt.h"
+
+extern "C" char* rust_bitcoin_psbt(uint8_t *data, size_t len);
+
+bool PSBTCore(Span<const uint8_t> buffer)
+{
+    PartiallySignedTransaction psbt_mut;
+    std::string error;
+    if (!DecodeRawPSBT(psbt_mut, MakeByteSpan(buffer), error)) {
+        return false;
+    }
+    return true;
+}
+
+
+void Psbt(FuzzedDataProvider& provider)
+{
+    std::vector<uint8_t> buffer{provider.ConsumeRemainingBytes<uint8_t>()};
+    bool core{PSBTCore(buffer)};
+    std::string rust_bitcoin{rust_bitcoin_psbt(buffer.data(), buffer.size())};
+    if (rust_bitcoin == "invalid xonly public key" || 
+        rust_bitcoin == "bitcoin consensus encoding error") return;
+    if (core) assert(rust_bitcoin == "1");
+    else assert(rust_bitcoin == "0");
+}
\ No newline at end of file
diff --git a/targets/psbt.h b/targets/psbt.h
new file mode 100644
index 0000000..4ae2572
--- /dev/null
+++ b/targets/psbt.h
@@ -0,0 +1,7 @@
+#ifndef PSBT_H
+#define PSBT_H
+
+#include <fuzzer/FuzzedDataProvider.h>
+
+void Psbt(FuzzedDataProvider& provider);
+#endif
\ No newline at end of file