From b67603827b24272517b43a7249b0f63625308735 Mon Sep 17 00:00:00 2001
From: Joey Sacchini <joey@sacchini.net>
Date: Fri, 8 Jan 2021 15:37:01 -0500
Subject: add some functions to ensure buffer is allocated on a CraftIO
 implementor

---
 src/connection.rs | 11 +++++++++++
 src/reader.rs     | 20 ++++++++++++++++++++
 src/wrapper.rs    |  9 +++++++--
 src/writer.rs     | 17 +++++++++++++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index c5e0cf7..94abbcd 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -46,6 +46,17 @@ impl<R, W> CraftIo for CraftConnection<R, W> {
         self.reader.set_max_packet_size(max_size);
         self.writer.set_max_packet_size(max_size);
     }
+
+    fn ensure_buf_capacity(&mut self, capacity: usize) {
+        self.reader.ensure_buf_capacity(capacity);
+        self.writer.ensure_buf_capacity(capacity);
+    }
+
+    #[cfg(feature = "compression")]
+    fn ensure_compression_buf_capacity(&mut self, capacity: usize) {
+        self.reader.ensure_buf_capacity(capacity);
+        self.writer.ensure_buf_capacity(capacity);
+    }
 }
 
 impl<R, W> CraftSyncReader for CraftConnection<R, W>
diff --git a/src/reader.rs b/src/reader.rs
index bba028e..b3bb2dd 100644
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -172,6 +172,26 @@ impl<R> CraftIo for CraftReader<R> {
         debug_assert!(max_size > 5);
         self.max_packet_size = max_size;
     }
+
+    fn ensure_buf_capacity(&mut self, capacity: usize) {
+        let alloc_to = if capacity > self.max_packet_size {
+            self.max_packet_size
+        } else {
+            capacity
+        };
+        self.move_ready_data_to_front();
+        get_sized_buf(&mut self.raw_buf, 0, alloc_to);
+    }
+
+    #[cfg(feature = "compression")]
+    fn ensure_compression_buf_capacity(&mut self, capacity: usize) {
+        let alloc_to = if capacity > self.max_packet_size {
+            self.max_packet_size
+        } else {
+            capacity
+        };
+        get_sized_buf(&mut self.decompress_buf, 0, alloc_to);
+    }
 }
 
 macro_rules! rr_unwrap {
diff --git a/src/wrapper.rs b/src/wrapper.rs
index 19ec840..cd21864 100644
--- a/src/wrapper.rs
+++ b/src/wrapper.rs
@@ -28,7 +28,6 @@ pub trait CraftIo {
     fn set_state(&mut self, next: State);
 
 
-    #[cfg(feature = "compression")]
     ///
     /// Modifies the compression configuration. If a value of `None` is provided, then compression is
     /// disabled. If a value of `Some` is provided, then the threshold is set to that value.
@@ -36,9 +35,9 @@ pub trait CraftIo {
     /// If a 0 or negative value is provided in a `Some` variant, then it is the same as calling
     /// this function with the `None` variant
     ///
+    #[cfg(feature = "compression")]
     fn set_compression_threshold(&mut self, threshold: Option<i32>);
 
-    #[cfg(feature = "encryption")]
     ///
     /// Modifies the encryption configuration. This function should only be called once, and can only
     /// be used to enable encryption.
@@ -46,6 +45,7 @@ pub trait CraftIo {
     /// If encryption is already enabled or the arguments are not valid for the cipher, then an
     /// error is returned and nothing in the underlying state is changed.
     ///
+    #[cfg(feature = "encryption")]
     fn enable_encryption(&mut self, key: &[u8], iv: &[u8]) -> Result<(), CipherError>;
 
     ///
@@ -60,4 +60,9 @@ pub trait CraftIo {
     /// todo split the compressed vs not compressed limits?
     ///
     fn set_max_packet_size(&mut self, max_size: usize);
+
+    fn ensure_buf_capacity(&mut self, capacity: usize);
+
+    #[cfg(feature = "compression")]
+    fn ensure_compression_buf_capacity(&mut self, capacity: usize);
 }
diff --git a/src/writer.rs b/src/writer.rs
index 4c0a302..0be6133 100644
--- a/src/writer.rs
+++ b/src/writer.rs
@@ -216,6 +216,23 @@ impl<W> CraftIo for CraftWriter<W> {
         debug_assert!(max_size > 5);
         self.max_packet_size = max_size;
     }
+
+    fn ensure_buf_capacity(&mut self, capacity: usize) {
+        get_sized_buf(&mut self.raw_buf, 0, if capacity > self.max_packet_size {
+            self.max_packet_size
+        } else {
+            capacity
+        });
+    }
+
+    #[cfg(feature = "compression")]
+    fn ensure_compression_buf_capacity(&mut self, capacity: usize) {
+        get_sized_buf(&mut self.compress_buf, 0, if capacity > self.max_packet_size {
+            self.max_packet_size
+        } else {
+            capacity
+        });
+    }
 }
 
 impl<W> CraftSyncWriter for CraftWriter<W>
-- 
cgit