From 1b8f8fd8da08823d88108a098b96abe591555eb9 Mon Sep 17 00:00:00 2001
From: b1n <b1n@b1n.io>
Date: Mon, 15 Jul 2024 20:17:39 +0800
Subject: [PATCH] fix(replace): no error returned when holder not found

---
 Cargo.lock    |  2 +-
 Cargo.toml    |  2 +-
 build.rs      | 18 ++++++++----------
 src/plugin.rs |  6 +++---
 src/utils.rs  | 28 ++++++++++++++++++----------
 5 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3de3dff..17c663f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3177,7 +3177,7 @@ dependencies = [
 
 [[package]]
 name = "pumpbin"
-version = "1.1.0"
+version = "1.1.1"
 dependencies = [
  "anyhow",
  "bincode",
diff --git a/Cargo.toml b/Cargo.toml
index c23ed51..0abec1b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pumpbin"
-version = "1.1.0"
+version = "1.1.1"
 authors = ["b1n <b1n@b1n.io>"]
 edition = "2021"
 description = "PumpBin is an Implant Generation Platform."
diff --git a/build.rs b/build.rs
index e154a9e..1d0020d 100644
--- a/build.rs
+++ b/build.rs
@@ -1,15 +1,13 @@
-fn build_capnp() {
-    capnpc::CompilerCommand::new()
-        .src_prefix("capnp")
-        .file("capnp/plugin.capnp")
-        .output_path("capnp")
-        .run()
-        .expect("schema compiler command");
-}
-
 fn main() {
     #[cfg(debug_assertions)]
-    build_capnp();
+    {
+        capnpc::CompilerCommand::new()
+            .src_prefix("capnp")
+            .file("capnp/plugin.capnp")
+            .output_path("capnp")
+            .run()
+            .expect("schema compiler command");
+    }
 
     #[cfg(target_os = "windows")]
     {
diff --git a/src/plugin.rs b/src/plugin.rs
index b44e521..309a2bd 100644
--- a/src/plugin.rs
+++ b/src/plugin.rs
@@ -479,14 +479,14 @@ impl Plugin {
             self.replace().src_prefix(),
             shellcode_src.as_slice(),
             self.replace().max_len(),
-        );
+        )?;
 
         // replace pass
         for pass in pass {
             let holder = pass.holder();
             let replace_by = pass.replace_by();
 
-            utils::replace(bin, holder, replace_by, holder.len());
+            utils::replace(bin, holder, replace_by, holder.len())?;
         }
 
         // replace size_holder
@@ -503,7 +503,7 @@ impl Plugin {
                 .collect();
             size_bytes.extend_from_slice(shellcode_len_bytes.as_slice());
 
-            utils::replace(bin, size_holder, size_bytes.as_slice(), size_holder.len());
+            utils::replace(bin, size_holder, size_bytes.as_slice(), size_holder.len())?;
         }
 
         Ok(())
diff --git a/src/utils.rs b/src/utils.rs
index 4044329..f934efd 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,5 +1,6 @@
 use std::iter;
 
+use anyhow::anyhow;
 use iced::{
     advanced::graphics::image::image_rs::ImageFormat,
     window::{self, Level, Position},
@@ -62,17 +63,24 @@ pub fn window_settings() -> window::Settings {
     }
 }
 
-pub fn replace(bin: &mut [u8], holder: &[u8], replace_by: &[u8], max_len: usize) {
+pub fn replace(
+    bin: &mut [u8],
+    holder: &[u8],
+    replace_by: &[u8],
+    max_len: usize,
+) -> anyhow::Result<()> {
     let mut replace_by = replace_by.to_owned();
 
-    let find = memmem::find_iter(bin, holder).next();
-    if let Some(position) = find {
-        let mut random: Vec<u8> = iter::repeat(b'0')
-            .take(max_len - replace_by.len())
-            .collect();
-        rand::thread_rng().fill_bytes(&mut random);
-        replace_by.extend_from_slice(random.as_slice());
+    let position = memmem::find_iter(bin, holder)
+        .next()
+        .ok_or(anyhow!("Not found {}", String::from_utf8_lossy(holder)))?;
+    let mut random: Vec<u8> = iter::repeat(b'0')
+        .take(max_len - replace_by.len())
+        .collect();
+    rand::thread_rng().fill_bytes(&mut random);
+    replace_by.extend_from_slice(random.as_slice());
 
-        bin[position..(position + max_len)].copy_from_slice(replace_by.as_slice());
-    }
+    bin[position..(position + max_len)].copy_from_slice(replace_by.as_slice());
+
+    Ok(())
 }