netget/server/netconf/
owned_stream.rs

1//! Cancellation reaches russh's internal driver without duplicating the socket.
2use std::{
3    future::Future,
4    io,
5    pin::Pin,
6    task::{Context, Poll},
7};
8use tokio::{
9    io::{AsyncRead, AsyncWrite, ReadBuf},
10    net::TcpStream,
11};
12use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
13
14pub struct Owner(CancellationToken);
15impl Owner {
16    pub fn token(&self) -> CancellationToken {
17        self.0.clone()
18    }
19}
20impl Drop for Owner {
21    fn drop(&mut self) {
22        self.0.cancel();
23    }
24}
25pub struct OwnedStream {
26    inner: TcpStream,
27    cancelled: Pin<Box<WaitForCancellationFutureOwned>>,
28}
29impl OwnedStream {
30    pub fn new(inner: TcpStream) -> (Self, Owner) {
31        let token = CancellationToken::new();
32        (
33            Self {
34                inner,
35                cancelled: Box::pin(token.clone().cancelled_owned()),
36            },
37            Owner(token),
38        )
39    }
40    fn cancellation(&mut self, cx: &mut Context<'_>) -> Option<io::Error> {
41        // Polling registers the internal driver task's waker even if the native
42        // socket stays silent. Owner drop wakes it and turns I/O into an error.
43        self.cancelled.as_mut().poll(cx).is_ready().then(|| {
44            io::Error::new(
45                io::ErrorKind::ConnectionAborted,
46                "NETCONF owned I/O cancelled",
47            )
48        })
49    }
50}
51impl AsyncRead for OwnedStream {
52    fn poll_read(
53        mut self: Pin<&mut Self>,
54        cx: &mut Context<'_>,
55        buffer: &mut ReadBuf<'_>,
56    ) -> Poll<io::Result<()>> {
57        if let Some(error) = self.cancellation(cx) {
58            return Poll::Ready(Err(error));
59        }
60        Pin::new(&mut self.inner).poll_read(cx, buffer)
61    }
62}
63impl AsyncWrite for OwnedStream {
64    fn poll_write(
65        mut self: Pin<&mut Self>,
66        cx: &mut Context<'_>,
67        bytes: &[u8],
68    ) -> Poll<io::Result<usize>> {
69        if let Some(error) = self.cancellation(cx) {
70            return Poll::Ready(Err(error));
71        }
72        Pin::new(&mut self.inner).poll_write(cx, bytes)
73    }
74    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
75        if let Some(error) = self.cancellation(cx) {
76            return Poll::Ready(Err(error));
77        }
78        Pin::new(&mut self.inner).poll_flush(cx)
79    }
80    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
81        if let Some(error) = self.cancellation(cx) {
82            return Poll::Ready(Err(error));
83        }
84        Pin::new(&mut self.inner).poll_shutdown(cx)
85    }
86}
netget/server/netconf/owned_stream.rs

netget/server/netconf/
owned_stream.rs
