/rust/registry/src/index.crates.io-1949cf8c6b5b557f/tokio-util-0.7.20/src/io/inspect.rs

Source
use pin_project_lite::pin_project;
use std::io::{IoSlice, Result};
use std::pin::Pin;
use std::task::{ready, Context, Poll};

use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

pin_project! {
    /// An adapter that lets you inspect the data that's being read.
    ///
    /// This is useful for things like hashing data as it's read in.
    pub struct InspectReader<R, F> {
        #[pin]
        reader: R,
        f: F,
    }
}

impl<R, F> InspectReader<R, F> {
    /// Create a new `InspectReader`, wrapping `reader` and calling `f` for the
    /// new data supplied by each read call.
    ///
    /// The closure will only be called with an empty slice if the inner reader
    /// returns without reading data into the buffer. This happens at EOF, or if
    /// `poll_read` is called with a zero-size buffer.
    pub fn new(reader: R, f: F) -> InspectReader<R, F>
    where
        R: AsyncRead,
        F: FnMut(&[u8]),
    {
        InspectReader { reader, f }
    }

    /// Acquires a reference to the underlying reader.
    pub fn get_ref(&self) -> &R {
        &self.reader
    }

    /// Acquires a mutable reference to the underlying reader.
    ///
    /// I/O performed directly on the underlying reader is not inspected.
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.reader
    }

    /// Acquires a pinned mutable reference to the underlying reader.
    ///
    /// I/O performed directly on the underlying reader is not inspected.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> {
        self.project().reader
    }

    /// Consumes the `InspectReader`, returning the wrapped reader
    pub fn into_inner(self) -> R {
        self.reader
    }
}

impl<R: AsyncRead, F: FnMut(&[u8])> AsyncRead for InspectReader<R, F> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<()>> {
        let me = self.project();
        let filled_length = buf.filled().len();
        ready!(me.reader.poll_read(cx, buf))?;
        (me.f)(&buf.filled()[filled_length..]);
        Poll::Ready(Ok(()))
    }
}

impl<R: AsyncWrite, F> AsyncWrite for InspectReader<R, F> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::result::Result<usize, std::io::Error>> {
        self.project().reader.poll_write(cx, buf)
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<std::result::Result<(), std::io::Error>> {
        self.project().reader.poll_flush(cx)
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<std::result::Result<(), std::io::Error>> {
        self.project().reader.poll_shutdown(cx)
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<Result<usize>> {
        self.project().reader.poll_write_vectored(cx, bufs)
    }

    fn is_write_vectored(&self) -> bool {
        self.reader.is_write_vectored()
    }
}

pin_project! {
    /// An adapter that lets you inspect the data that's being written.
    ///
    /// This is useful for things like hashing data as it's written out.
    pub struct InspectWriter<W, F> {
        #[pin]
        writer: W,
        f: F,
    }
}

impl<W, F> InspectWriter<W, F> {
    /// Create a new `InspectWriter`, wrapping `write` and calling `f` for the
    /// data successfully written by each write call.
    ///
    /// The closure `f` will never be called with an empty slice. A vectored
    /// write can result in multiple calls to `f` - at most one call to `f` per
    /// buffer supplied to `poll_write_vectored`.
    pub fn new(writer: W, f: F) -> InspectWriter<W, F>
    where
        W: AsyncWrite,
        F: FnMut(&[u8]),
    {
        InspectWriter { writer, f }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.writer
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// I/O performed directly on the underlying writer is not inspected.
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Acquires a pinned mutable reference to the underlying writer.
    ///
    /// I/O performed directly on the underlying writer is not inspected.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> {
        self.project().writer
    }

    /// Consumes the `InspectWriter`, returning the wrapped writer
    pub fn into_inner(self) -> W {
        self.writer
    }
}

impl<W: AsyncWrite, F: FnMut(&[u8])> AsyncWrite for InspectWriter<W, F> {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
        let me = self.project();
        let res = me.writer.poll_write(cx, buf);
        if let Poll::Ready(Ok(count)) = res {
            if count != 0 {
                (me.f)(&buf[..count]);
            }
        }
        res
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        let me = self.project();
        me.writer.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        let me = self.project();
        me.writer.poll_shutdown(cx)
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<Result<usize>> {
        let me = self.project();
        let res = me.writer.poll_write_vectored(cx, bufs);
        if let Poll::Ready(Ok(mut count)) = res {
            for buf in bufs {
                if count == 0 {
                    break;
                }
                let size = count.min(buf.len());
                if size != 0 {
                    (me.f)(&buf[..size]);
                    count -= size;
                }
            }
        }
        res
    }

    fn is_write_vectored(&self) -> bool {
        self.writer.is_write_vectored()
    }
}

impl<W: AsyncRead, F> AsyncRead for InspectWriter<W, F> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        self.project().writer.poll_read(cx, buf)
    }
}

Coverage Report

Created: 2026-10-17 06:21

Line	Count	Source
1		use pin_project_lite::pin_project;
2		use std::io::{IoSlice, Result};
3		use std::pin::Pin;
4		use std::task::{ready, Context, Poll};
5
6		use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
7
8		pin_project! {
9		/// An adapter that lets you inspect the data that's being read.
10		///
11		/// This is useful for things like hashing data as it's read in.
12		pub struct InspectReader<R, F> {
13		#[pin]
14		reader: R,
15		f: F,
16		}
17		}
18
19		impl<R, F> InspectReader<R, F> {
20		/// Create a new `InspectReader`, wrapping `reader` and calling `f` for the
21		/// new data supplied by each read call.
22		///
23		/// The closure will only be called with an empty slice if the inner reader
24		/// returns without reading data into the buffer. This happens at EOF, or if
25		/// `poll_read` is called with a zero-size buffer.
26	0	pub fn new(reader: R, f: F) -> InspectReader<R, F>
27	0	where
28	0	R: AsyncRead,
29	0	F: FnMut(&[u8]),
30		{
31	0	InspectReader { reader, f }
32	0	}
33
34		/// Acquires a reference to the underlying reader.
35	0	pub fn get_ref(&self) -> &R {
36	0	&self.reader
37	0	}
38
39		/// Acquires a mutable reference to the underlying reader.
40		///
41		/// I/O performed directly on the underlying reader is not inspected.
42	0	pub fn get_mut(&mut self) -> &mut R {
43	0	&mut self.reader
44	0	}
45
46		/// Acquires a pinned mutable reference to the underlying reader.
47		///
48		/// I/O performed directly on the underlying reader is not inspected.
49	0	pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> {
50	0	self.project().reader
51	0	}
52
53		/// Consumes the `InspectReader`, returning the wrapped reader
54	0	pub fn into_inner(self) -> R {
55	0	self.reader
56	0	}
57		}
58
59		impl<R: AsyncRead, F: FnMut(&[u8])> AsyncRead for InspectReader<R, F> {
60	0	fn poll_read(
61	0	self: Pin<&mut Self>,
62	0	cx: &mut Context<'_>,
63	0	buf: &mut ReadBuf<'_>,
64	0	) -> Poll<Result<()>> {
65	0	let me = self.project();
66	0	let filled_length = buf.filled().len();
67	0	ready!(me.reader.poll_read(cx, buf))?;
68	0	(me.f)(&buf.filled()[filled_length..]);
69	0	Poll::Ready(Ok(()))
70	0	}
71		}
72
73		impl<R: AsyncWrite, F> AsyncWrite for InspectReader<R, F> {
74	0	fn poll_write(
75	0	self: Pin<&mut Self>,
76	0	cx: &mut Context<'_>,
77	0	buf: &[u8],
78	0	) -> Poll<std::result::Result<usize, std::io::Error>> {
79	0	self.project().reader.poll_write(cx, buf)
80	0	}
81
82	0	fn poll_flush(
83	0	self: Pin<&mut Self>,
84	0	cx: &mut Context<'_>,
85	0	) -> Poll<std::result::Result<(), std::io::Error>> {
86	0	self.project().reader.poll_flush(cx)
87	0	}
88
89	0	fn poll_shutdown(
90	0	self: Pin<&mut Self>,
91	0	cx: &mut Context<'_>,
92	0	) -> Poll<std::result::Result<(), std::io::Error>> {
93	0	self.project().reader.poll_shutdown(cx)
94	0	}
95
96	0	fn poll_write_vectored(
97	0	self: Pin<&mut Self>,
98	0	cx: &mut Context<'_>,
99	0	bufs: &[IoSlice<'_>],
100	0	) -> Poll<Result<usize>> {
101	0	self.project().reader.poll_write_vectored(cx, bufs)
102	0	}
103
104	0	fn is_write_vectored(&self) -> bool {
105	0	self.reader.is_write_vectored()
106	0	}
107		}
108
109		pin_project! {
110		/// An adapter that lets you inspect the data that's being written.
111		///
112		/// This is useful for things like hashing data as it's written out.
113		pub struct InspectWriter<W, F> {
114		#[pin]
115		writer: W,
116		f: F,
117		}
118		}
119
120		impl<W, F> InspectWriter<W, F> {
121		/// Create a new `InspectWriter`, wrapping `write` and calling `f` for the
122		/// data successfully written by each write call.
123		///
124		/// The closure `f` will never be called with an empty slice. A vectored
125		/// write can result in multiple calls to `f` - at most one call to `f` per
126		/// buffer supplied to `poll_write_vectored`.
127	0	pub fn new(writer: W, f: F) -> InspectWriter<W, F>
128	0	where
129	0	W: AsyncWrite,
130	0	F: FnMut(&[u8]),
131		{
132	0	InspectWriter { writer, f }
133	0	}
134
135		/// Acquires a reference to the underlying writer.
136	0	pub fn get_ref(&self) -> &W {
137	0	&self.writer
138	0	}
139
140		/// Acquires a mutable reference to the underlying writer.
141		///
142		/// I/O performed directly on the underlying writer is not inspected.
143	0	pub fn get_mut(&mut self) -> &mut W {
144	0	&mut self.writer
145	0	}
146
147		/// Acquires a pinned mutable reference to the underlying writer.
148		///
149		/// I/O performed directly on the underlying writer is not inspected.
150	0	pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> {
151	0	self.project().writer
152	0	}
153
154		/// Consumes the `InspectWriter`, returning the wrapped writer
155	0	pub fn into_inner(self) -> W {
156	0	self.writer
157	0	}
158		}
159
160		impl<W: AsyncWrite, F: FnMut(&[u8])> AsyncWrite for InspectWriter<W, F> {
161	0	fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
162	0	let me = self.project();
163	0	let res = me.writer.poll_write(cx, buf);
164	0	if let Poll::Ready(Ok(count)) = res {
165	0	if count != 0 {
166	0	(me.f)(&buf[..count]);
167	0	}
168	0	}
169	0	res
170	0	}
171
172	0	fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
173	0	let me = self.project();
174	0	me.writer.poll_flush(cx)
175	0	}
176
177	0	fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
178	0	let me = self.project();
179	0	me.writer.poll_shutdown(cx)
180	0	}
181
182	0	fn poll_write_vectored(
183	0	self: Pin<&mut Self>,
184	0	cx: &mut Context<'_>,
185	0	bufs: &[IoSlice<'_>],
186	0	) -> Poll<Result<usize>> {
187	0	let me = self.project();
188	0	let res = me.writer.poll_write_vectored(cx, bufs);
189	0	if let Poll::Ready(Ok(mut count)) = res {
190	0	for buf in bufs {
191	0	if count == 0 {
192	0	break;
193	0	}
194	0	let size = count.min(buf.len());
195	0	if size != 0 {
196	0	(me.f)(&buf[..size]);
197	0	count -= size;
198	0	}
199		}
200	0	}
201	0	res
202	0	}
203
204	0	fn is_write_vectored(&self) -> bool {
205	0	self.writer.is_write_vectored()
206	0	}
207		}
208
209		impl<W: AsyncRead, F> AsyncRead for InspectWriter<W, F> {
210	0	fn poll_read(
211	0	self: Pin<&mut Self>,
212	0	cx: &mut Context<'_>,
213	0	buf: &mut ReadBuf<'_>,
214	0	) -> Poll<std::io::Result<()>> {
215	0	self.project().writer.poll_read(cx, buf)
216	0	}
217		}