From d0515dc6eff82aff43ff02d26921d6f48ac4cbdb Mon Sep 17 00:00:00 2001
From: Dylan Plecki <dylan@plecki.net>
Date: Sat, 27 Jan 2024 15:04:25 -0500
Subject: [PATCH] Add SeekBuf trait impl to Take

---
 src/buf/take.rs    | 22 +++++++++++++++++++++-
 tests/test_take.rs | 15 ++++++++++++++-
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/buf/take.rs b/src/buf/take.rs
index a16a434ee..c2cfa9bc3 100644
--- a/src/buf/take.rs
+++ b/src/buf/take.rs
@@ -1,4 +1,4 @@
-use crate::{Buf, Bytes};
+use crate::{Buf, Bytes, SeekBuf};
 
 use core::cmp;
 
@@ -153,3 +153,23 @@ impl<T: Buf> Buf for Take<T> {
         r
     }
 }
+
+impl<T: SeekBuf> SeekBuf for Take<T> {
+    fn chunk_from(&self, start: usize) -> Option<&[u8]> {
+        if start < self.limit {
+            let remaining = self.limit - start;
+            let chunk = self.inner.chunk_from(start)?;
+            Some(&chunk[..chunk.len().min(remaining)])
+        } else {
+            None
+        }
+    }
+
+    fn chunk_to(&self, end: usize) -> Option<&[u8]> {
+        if end <= self.limit {
+            self.inner.chunk_to(end)
+        } else {
+            None
+        }
+    }
+}
diff --git a/tests/test_take.rs b/tests/test_take.rs
index 51df91d14..6768f4301 100644
--- a/tests/test_take.rs
+++ b/tests/test_take.rs
@@ -1,7 +1,7 @@
 #![warn(rust_2018_idioms)]
 
 use bytes::buf::Buf;
-use bytes::Bytes;
+use bytes::{Bytes, SeekBuf};
 
 #[test]
 fn long_take() {
@@ -12,6 +12,19 @@ fn long_take() {
     assert_eq!(b"hello world", buf.chunk());
 }
 
+#[test]
+fn take_from_seek_buf() {
+    let buf = b"hello world".take(5);
+
+    assert_eq!(buf.chunk_from(0), Some(b"hello".as_slice()));
+    assert_eq!(buf.chunk_from(3), Some(b"lo".as_slice()));
+    assert_eq!(buf.chunk_to(5), Some(b"hello".as_slice()));
+    assert_eq!(buf.chunk_to(2), Some(b"he".as_slice()));
+
+    assert_eq!(buf.chunk_from(6), None);
+    assert_eq!(buf.chunk_to(6), None);
+}
+
 #[test]
 fn take_copy_to_bytes() {
     let mut abcd = Bytes::copy_from_slice(b"abcd");