diff --git a/yazi-sftp/src/error.rs b/yazi-sftp/src/error.rs index 1fd74de9..0fe10a89 100644 --- a/yazi-sftp/src/error.rs +++ b/yazi-sftp/src/error.rs @@ -59,6 +59,10 @@ impl From for Error { fn from(_: tokio::sync::oneshot::error::RecvError) -> Self { Self::custom("channel closed") } } +impl From for Error { + fn from(_: tokio::time::error::Elapsed) -> Self { Self::Timeout } +} + impl std::error::Error for Error {} impl std::fmt::Display for Error { diff --git a/yazi-sftp/src/fs/file.rs b/yazi-sftp/src/fs/file.rs index fca67713..71b53c77 100644 --- a/yazi-sftp/src/fs/file.rs +++ b/yazi-sftp/src/fs/file.rs @@ -1,8 +1,8 @@ use std::{io, pin::Pin, sync::Arc, task::{Context, Poll, ready}, time::Duration}; -use tokio::{io::{AsyncRead, AsyncWrite, ReadBuf}, sync::oneshot, time::{Timeout, timeout}}; +use tokio::{io::{AsyncRead, AsyncWrite, ReadBuf}, time::{Timeout, timeout}}; -use crate::{Error, Operator, Packet, Session, fs::Attrs}; +use crate::{Error, Operator, Packet, Receiver, Session, fs::Attrs}; pub struct File { session: Arc, @@ -10,10 +10,10 @@ pub struct File { closed: bool, cursor: u64, - close_rx: Option>>>, - read_rx: Option>>, - write_rx: Option<(oneshot::Receiver>, usize)>, - flush_rx: Option>>>, + close_rx: Option>, + read_rx: Option, + write_rx: Option<(Receiver, usize)>, + flush_rx: Option>, } impl Unpin for File {} diff --git a/yazi-sftp/src/lib.rs b/yazi-sftp/src/lib.rs index 2b6d05c0..ea980de9 100644 --- a/yazi-sftp/src/lib.rs +++ b/yazi-sftp/src/lib.rs @@ -9,6 +9,7 @@ mod macros; mod operator; mod packet; mod path; +mod receiver; mod ser; mod session; @@ -18,5 +19,6 @@ pub(crate) use id::*; pub use operator::*; pub use packet::*; pub use path::*; +pub use receiver::*; pub(crate) use ser::*; pub use session::*; diff --git a/yazi-sftp/src/operator.rs b/yazi-sftp/src/operator.rs index b28cb181..143e1963 100644 --- a/yazi-sftp/src/operator.rs +++ b/yazi-sftp/src/operator.rs @@ -1,15 +1,14 @@ use std::{ops::Deref, sync::Arc}; use russh::{ChannelStream, client::Msg}; -use tokio::sync::oneshot; use typed_path::UnixPathBuf; -use crate::{AsSftpPath, Error, Packet, Session, SftpPath, fs::{Attrs, File, Flags, ReadDir}, requests, responses}; +use crate::{AsSftpPath, Error, Receiver, Session, SftpPath, fs::{Attrs, File, Flags, ReadDir}, requests, responses}; pub struct Operator(Arc); impl Deref for Operator { - type Target = Session; + type Target = Arc; fn deref(&self) -> &Self::Target { &self.0 } } @@ -36,25 +35,15 @@ impl Operator { Ok(File::new(&self.0, handle.handle)) } - pub fn close(&self, handle: &str) -> Result>, Error> { + pub fn close(&self, handle: &str) -> Result { self.send_sync(requests::Close::new(handle)) } - pub fn read( - &self, - handle: &str, - offset: u64, - len: u32, - ) -> Result>, Error> { + pub fn read(&self, handle: &str, offset: u64, len: u32) -> Result { self.send_sync(requests::Read::new(handle, offset, len)) } - pub fn write( - &self, - handle: &str, - offset: u64, - data: &[u8], - ) -> Result>, Error> { + pub fn write(&self, handle: &str, offset: u64, data: &[u8]) -> Result { self.send_sync(requests::Write::new(handle, offset, data)) } @@ -183,7 +172,7 @@ impl Operator { status.into() } - pub fn fsync(&self, handle: &str) -> Result>, Error> { + pub fn fsync(&self, handle: &str) -> Result { if self.extensions.lock().get("fsync@openssh.com").is_none_or(|s| s != "1") { return Err(Error::Unsupported); } diff --git a/yazi-sftp/src/receiver.rs b/yazi-sftp/src/receiver.rs new file mode 100644 index 00000000..c0304586 --- /dev/null +++ b/yazi-sftp/src/receiver.rs @@ -0,0 +1,47 @@ +use std::{pin::Pin, sync::Arc, task::Poll}; + +use tokio::sync::oneshot; + +use crate::{Packet, Session}; + +pub struct Receiver { + rx: oneshot::Receiver>, + received: bool, + + session: Arc, + id: u32, +} + +impl Drop for Receiver { + fn drop(&mut self) { + if !self.received { + self.session.callback.lock().remove(&self.id); + } + } +} + +impl Receiver { + pub(crate) fn new( + session: &Arc, + id: u32, + rx: oneshot::Receiver>, + ) -> Self { + Self { rx, received: false, session: session.clone(), id } + } +} + +impl Future for Receiver { + type Output = Result, oneshot::error::RecvError>; + + fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { + let me = self.get_mut(); + match Pin::new(&mut me.rx).poll(cx) { + Poll::Ready(Ok(packet)) => { + me.received = true; + Poll::Ready(Ok(packet)) + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } +} diff --git a/yazi-sftp/src/responses/status.rs b/yazi-sftp/src/responses/status.rs index fc864c89..137267df 100644 --- a/yazi-sftp/src/responses/status.rs +++ b/yazi-sftp/src/responses/status.rs @@ -27,6 +27,15 @@ impl Status { pub fn is_ok(&self) -> bool { self.code == StatusCode::Ok } pub fn is_eof(&self) -> bool { self.code == StatusCode::Eof } + + pub(crate) fn connection_lost(id: u32) -> Self { + Self { + id, + code: StatusCode::ConnectionLost, + message: "connection lost".to_owned(), + language: "en".to_owned(), + } + } } #[derive(Debug, Clone, Copy, PartialEq, Deserialize, Serialize)] diff --git a/yazi-sftp/src/session.rs b/yazi-sftp/src/session.rs index 1f1c1c44..0857bc2c 100644 --- a/yazi-sftp/src/session.rs +++ b/yazi-sftp/src/session.rs @@ -1,16 +1,16 @@ -use std::{any::TypeId, collections::HashMap, io::{self, ErrorKind}, sync::Arc}; +use std::{any::TypeId, collections::HashMap, io::{self, ErrorKind}, sync::Arc, time::Duration}; use parking_lot::Mutex; use russh::{ChannelStream, client::Msg}; use serde::Serialize; -use tokio::{io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf}, select, sync::{mpsc, oneshot}}; +use tokio::{io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf}, select, sync::{mpsc, oneshot}, time::timeout}; -use crate::{Error, Id, Packet, responses}; +use crate::{Error, Id, Packet, Receiver, responses}; pub struct Session { tx: mpsc::UnboundedSender>, id: Id, - callback: Mutex>>>, + pub(super) callback: Mutex>>>, pub(super) extensions: Mutex>, } @@ -43,6 +43,7 @@ impl Session { } } + let me_ = me.clone(); let (mut reader, mut writer) = tokio::io::split(stream); tokio::spawn(async move { while let Some(data) = rx.recv().await { @@ -51,6 +52,9 @@ impl Session { { rx.close(); writer.shutdown().await.ok(); + for (id, cb) in me_.callback.lock().drain() { + cb.send(responses::Status::connection_lost(id).into()).ok(); + } break; } } @@ -84,12 +88,12 @@ impl Session { me } - pub async fn send<'a, I, O>(&self, input: I) -> Result + pub async fn send<'a, I, O>(self: &Arc, input: I) -> Result where I: Into> + Serialize, O: TryFrom, Error = Error> + 'static, { - match self.send_sync(input)?.await? { + match timeout(Duration::from_secs(30), self.send_sync(input)?).await?? { Packet::Status(status) if TypeId::of::() != TypeId::of::() => { Err(Error::Status(status)) } @@ -97,7 +101,7 @@ impl Session { } } - pub fn send_sync<'a, I>(&self, input: I) -> Result>, Error> + pub fn send_sync<'a, I>(self: &Arc, input: I) -> Result where I: Into> + Serialize, { @@ -106,12 +110,13 @@ impl Session { request = request.with_id(self.id.next()); } + let id = request.id(); let (tx, rx) = oneshot::channel(); - self.callback.lock().insert(request.id(), tx); - self.tx.send(crate::to_bytes(request)?)?; - Ok(rx) + self.callback.lock().insert(id, tx); + self.tx.send(crate::to_bytes(request)?)?; + Ok(Receiver::new(self, id, rx)) } - pub fn is_closed(&self) -> bool { self.tx.is_closed() } + pub fn is_closed(self: &Arc) -> bool { self.tx.is_closed() } }