Skip to main content

kube_portforward/
stream.rs

1use std::{
2    io,
3    pin::Pin,
4    task::{
5        Context,
6        Poll,
7    },
8};
9
10use tokio::io::{
11    AsyncBufRead,
12    AsyncRead,
13    AsyncWrite,
14    ReadBuf,
15};
16
17/// port forward stream backed by one stream pair on
18/// a multiplexed connection to the apiserver. Implements
19/// `AsyncRead + AsyncWrite` on the data half. Dropping the `Stream` (or
20/// calling `poll_shutdown`) sends a SPDY FIN frame so the apiserver tears
21/// down the backing pod connection right away.
22pub struct Stream {
23    inner: Box<spdy_mux::Stream>,
24}
25
26impl Stream {
27    pub(crate) fn from_spdy(stream: spdy_mux::Stream) -> Self {
28        Self {
29            inner: Box::new(stream),
30        }
31    }
32
33    /// Returns true if the remote already closed this stream's read
34    /// side.
35    pub fn is_read_closed(&self) -> bool {
36        self.inner.is_read_closed()
37    }
38
39    /// Split the stream into its data half (`AsyncRead + AsyncWrite`) and
40    /// error half (`AsyncRead`-only).
41    pub fn split(self) -> (DataStream, ErrorStream) {
42        let (d, e) = (*self.inner).split();
43        (DataStream { inner: d }, ErrorStream { inner: e })
44    }
45}
46
47impl AsyncRead for Stream {
48    fn poll_read(
49        self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>,
50    ) -> Poll<io::Result<()>> {
51        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
52    }
53}
54
55impl AsyncBufRead for Stream {
56    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
57        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
58    }
59
60    fn consume(self: Pin<&mut Self>, amt: usize) {
61        Pin::new(&mut self.get_mut().inner).consume(amt);
62    }
63}
64
65impl AsyncWrite for Stream {
66    fn poll_write(
67        self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8],
68    ) -> Poll<io::Result<usize>> {
69        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
70    }
71
72    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
73        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
74    }
75
76    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
77        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
78    }
79}
80
81pub struct DataStream {
82    inner: spdy_mux::DataStream,
83}
84
85impl AsyncRead for DataStream {
86    fn poll_read(
87        self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>,
88    ) -> Poll<io::Result<()>> {
89        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
90    }
91}
92
93impl AsyncBufRead for DataStream {
94    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
95        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
96    }
97
98    fn consume(self: Pin<&mut Self>, amt: usize) {
99        Pin::new(&mut self.get_mut().inner).consume(amt);
100    }
101}
102
103impl AsyncWrite for DataStream {
104    fn poll_write(
105        self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8],
106    ) -> Poll<io::Result<usize>> {
107        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
108    }
109
110    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
111        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
112    }
113
114    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
115        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
116    }
117}
118
119pub struct ErrorStream {
120    inner: spdy_mux::ErrorStream,
121}
122
123impl AsyncRead for ErrorStream {
124    fn poll_read(
125        self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>,
126    ) -> Poll<io::Result<()>> {
127        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
128    }
129}
130
131const _ASSERT_TRAITS: fn() = || {
132    const fn assert<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>() {}
133    assert::<Stream>();
134    assert::<DataStream>();
135};