use crate::{Buf, Bytes}; use core::cmp; /// A `Buf` adapter which limits the bytes read from an underlying buffer. /// /// This struct is generally created by calling `take()` on `Buf`. See /// documentation of [`take()`](trait.Buf.html#method.take) for more details. #[derive(Debug)] pub struct Take { inner: T, limit: usize, } pub fn new(inner: T, limit: usize) -> Take { Take { inner, limit } } impl Take { /// Consumes this `Take`, returning the underlying value. /// /// # Examples /// /// ```rust /// use bytes::{Buf, BufMut}; /// /// let mut buf = b"hello world".take(2); /// let mut dst = vec![]; /// /// dst.put(&mut buf); /// assert_eq!(*dst, b"he"[..]); /// /// let mut buf = buf.into_inner(); /// /// dst.clear(); /// dst.put(&mut buf); /// assert_eq!(*dst, b"llo world"[..]); /// ``` pub fn into_inner(self) -> T { self.inner } /// Gets a reference to the underlying `Buf`. /// /// It is inadvisable to directly read from the underlying `Buf`. /// /// # Examples /// /// ```rust /// use bytes::Buf; /// /// let buf = b"hello world".take(2); /// /// assert_eq!(11, buf.get_ref().remaining()); /// ``` pub fn get_ref(&self) -> &T { &self.inner } /// Gets a mutable reference to the underlying `Buf`. /// /// It is inadvisable to directly read from the underlying `Buf`. /// /// # Examples /// /// ```rust /// use bytes::{Buf, BufMut}; /// /// let mut buf = b"hello world".take(2); /// let mut dst = vec![]; /// /// buf.get_mut().advance(2); /// /// dst.put(&mut buf); /// assert_eq!(*dst, b"ll"[..]); /// ``` pub fn get_mut(&mut self) -> &mut T { &mut self.inner } /// Returns the maximum number of bytes that can be read. /// /// # Note /// /// If the inner `Buf` has fewer bytes than indicated by this method then /// that is the actual number of available bytes. /// /// # Examples /// /// ```rust /// use bytes::Buf; /// /// let mut buf = b"hello world".take(2); /// /// assert_eq!(2, buf.limit()); /// assert_eq!(b'h', buf.get_u8()); /// assert_eq!(1, buf.limit()); /// ``` pub fn limit(&self) -> usize { self.limit } /// Sets the maximum number of bytes that can be read. /// /// # Note /// /// If the inner `Buf` has fewer bytes than `lim` then that is the actual /// number of available bytes. /// /// # Examples /// /// ```rust /// use bytes::{Buf, BufMut}; /// /// let mut buf = b"hello world".take(2); /// let mut dst = vec![]; /// /// dst.put(&mut buf); /// assert_eq!(*dst, b"he"[..]); /// /// dst.clear(); /// /// buf.set_limit(3); /// dst.put(&mut buf); /// assert_eq!(*dst, b"llo"[..]); /// ``` pub fn set_limit(&mut self, lim: usize) { self.limit = lim } } impl Buf for Take { fn remaining(&self) -> usize { cmp::min(self.inner.remaining(), self.limit) } fn chunk(&self) -> &[u8] { let bytes = self.inner.chunk(); &bytes[..cmp::min(bytes.len(), self.limit)] } fn advance(&mut self, cnt: usize) { assert!(cnt <= self.limit); self.inner.advance(cnt); self.limit -= cnt; } fn copy_to_bytes(&mut self, len: usize) -> Bytes { assert!(len <= self.remaining(), "`len` greater than remaining"); let r = self.inner.copy_to_bytes(len); self.limit -= len; r } }