From 58ae8cc0474408cf1c19fa8db564b8329a56246a Mon Sep 17 00:00:00 2001
From: mzfr <github@mzfr.in>
Date: Sun, 29 Dec 2024 19:36:36 +0530
Subject: [PATCH] Fallback to to_allocvec in case of event_buffer overflow

Also combined the shared logic between compressed & uncompressed event
firing while keeping the same behavior
---
 libafl/src/events/llmp/mgr.rs | 84 +++++++++++++++--------------------
 1 file changed, 36 insertions(+), 48 deletions(-)

diff --git a/libafl/src/events/llmp/mgr.rs b/libafl/src/events/llmp/mgr.rs
index 7d3609b291..6426242473 100644
--- a/libafl/src/events/llmp/mgr.rs
+++ b/libafl/src/events/llmp/mgr.rs
@@ -518,68 +518,56 @@ where
             true
         }
     }
-
-    #[cfg(feature = "llmp_compression")]
     fn fire(
         &mut self,
         _state: &mut Self::State,
         event: Event<<Self::State as UsesInput>::Input>,
     ) -> Result<(), Error> {
+        #[cfg(feature = "llmp_compression")]
         let flags = LLMP_FLAG_INITIALIZED;
+
         self.event_buffer.clear();
         self.event_buffer.resize(self.event_buffer.capacity(), 0);
 
-        match postcard::to_slice(&event, &mut self.event_buffer) {
-            Ok(written) => {
-                let written_len = written.len();
-                match self
-                    .compressor
-                    .maybe_compress(&self.event_buffer[..written_len])
-                {
-                    Some(comp_buf) => {
-                        self.llmp.send_buf_with_flags(
-                            LLMP_TAG_EVENT_TO_BOTH,
-                            flags | LLMP_FLAG_COMPRESSED,
-                            &comp_buf,
-                        )?;
-                    }
-                    None => {
-                        self.llmp
-                            .send_buf(LLMP_TAG_EVENT_TO_BOTH, &self.event_buffer[..written_len])?;
-                    }
-                }
-                self.last_sent = current_time();
-                Ok(())
-            }
+        // Serialize the event, reallocating event_buffer if needed
+        let written_len = match postcard::to_slice(&event, &mut self.event_buffer) {
+            Ok(written) => written.len(),
             Err(postcard::Error::SerializeBufferFull) => {
-                return Err(Error::serialize("Buffer full"));
+                let serialized = postcard::to_allocvec(&event)?;
+                self.event_buffer = serialized;
+                self.event_buffer.len()
             }
-            Err(e) => Err(Error::from(e)),
-        }
-    }
-
-    #[cfg(not(feature = "llmp_compression"))]
-    fn fire(
-        &mut self,
-        _state: &mut Self::State,
-        event: Event<<Self::State as UsesInput>::Input>,
-    ) -> Result<(), Error> {
-        self.event_buffer.clear();
-        self.event_buffer.resize(self.event_buffer.capacity(), 0);
-
-        match postcard::to_slice(&event, &mut self.event_buffer) {
-            Ok(written) => {
-                let written_len = written.len();
-
-                self.llmp
-                    .send_buf(LLMP_TAG_EVENT_TO_BOTH, &self.event_buffer[..written_len])?;
+            Err(e) => return Err(Error::from(e)),
+        };
 
-                self.last_sent = current_time();
-                Ok(())
+        #[cfg(feature = "llmp_compression")]
+        {
+            match self
+                .compressor
+                .maybe_compress(&self.event_buffer[..written_len])
+            {
+                Some(comp_buf) => {
+                    self.llmp.send_buf_with_flags(
+                        LLMP_TAG_EVENT_TO_BOTH,
+                        flags | LLMP_FLAG_COMPRESSED,
+                        &comp_buf,
+                    )?;
+                }
+                None => {
+                    self.llmp
+                        .send_buf(LLMP_TAG_EVENT_TO_BOTH, &self.event_buffer[..written_len])?;
+                }
             }
-            Err(postcard::Error::SerializeBufferFull) => Err(Error::serialize("Buffer full")),
-            Err(e) => Err(Error::from(e)),
         }
+
+        #[cfg(not(feature = "llmp_compression"))]
+        {
+            self.llmp
+                .send_buf(LLMP_TAG_EVENT_TO_BOTH, &self.event_buffer[..written_len]);
+        }
+
+        self.last_sent = current_time();
+        Ok(())
     }
     fn serialize_observers<OT>(&mut self, observers: &OT) -> Result<Option<Vec<u8>>, Error>
     where