From eae2cf2cbf100c2c307ebb14aafc82d412c1581c Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 27 Jan 2023 12:11:45 +0100
Subject: [PATCH] Add IoRef::with_rw_buf() helper

---
 ntex-io/CHANGES.md   |  4 ++++
 ntex-io/Cargo.toml   |  2 +-
 ntex-io/src/buf.rs   | 14 ++++++++++++++
 ntex-io/src/ioref.rs | 13 +++++++++++--
 4 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/ntex-io/CHANGES.md b/ntex-io/CHANGES.md
index 720f01f2..e4f7968e 100644
--- a/ntex-io/CHANGES.md
+++ b/ntex-io/CHANGES.md
@@ -1,5 +1,9 @@
 # Changes
 
+## [0.2.6] - 2023-01-27
+
+* Add IoRef::with_rw_buf() helper
+
 ## [0.2.5] - 2023-01-27
 
 * Custom panic message for nested buffer borrow
diff --git a/ntex-io/Cargo.toml b/ntex-io/Cargo.toml
index 87fa483a..fc53c03e 100644
--- a/ntex-io/Cargo.toml
+++ b/ntex-io/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ntex-io"
-version = "0.2.5"
+version = "0.2.6"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Utilities for encoding and decoding frames"
 keywords = ["network", "framework", "async", "futures"]
diff --git a/ntex-io/src/buf.rs b/ntex-io/src/buf.rs
index 4ea3aef2..3c530da6 100644
--- a/ntex-io/src/buf.rs
+++ b/ntex-io/src/buf.rs
@@ -131,6 +131,20 @@ impl Stack {
             .unwrap_or(0)
     }
 
+    pub(crate) fn with_rw_buf<F, R>(&mut self, io: &IoRef, f: F) -> R
+    where
+        F: FnOnce(&mut BytesVec, &mut BytesVec) -> R,
+    {
+        let lvl = self.get_first_level();
+        if lvl.0.is_none() {
+            lvl.0 = Some(io.memory_pool().get_read_buf());
+        }
+        if lvl.1.is_none() {
+            lvl.1 = Some(io.memory_pool().get_write_buf());
+        }
+        f(lvl.0.as_mut().unwrap(), lvl.1.as_mut().unwrap())
+    }
+
     pub(crate) fn first_read_buf(&mut self) -> &mut Option<BytesVec> {
         &mut self.get_first_level().0
     }
diff --git a/ntex-io/src/ioref.rs b/ntex-io/src/ioref.rs
index b2a5fac1..492654f5 100644
--- a/ntex-io/src/ioref.rs
+++ b/ntex-io/src/ioref.rs
@@ -1,9 +1,9 @@
-use std::{any, cell, fmt, hash, io, time};
+use std::{any, fmt, hash, io, time, cell};
 
 use ntex_bytes::{BytesVec, PoolRef};
 use ntex_codec::{Decoder, Encoder};
 
-use super::{buf::Stack, io::Flags, timer, types, Filter, IoRef, OnDisconnect, WriteBuf};
+use super::{io::Flags, buf::Stack, timer, types, Filter, IoRef, OnDisconnect, WriteBuf};
 
 impl IoRef {
     #[inline]
@@ -198,6 +198,15 @@ impl IoRef {
         f(buf.as_mut().unwrap())
     }
 
+    #[inline]
+    /// Get mut access to read and write buffer
+    pub fn with_rw_buf<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut BytesVec, &mut BytesVec) -> R,
+    {
+        borrow_buffer(&self.0.buffer).with_rw_buf(self, f)
+    }
+
     #[inline]
     /// Start keep-alive timer
     pub fn start_keepalive_timer(&self, timeout: time::Duration) {