51 lines
1.3 KiB
Rust
51 lines
1.3 KiB
Rust
use russh::ChannelId;
|
|
use russh::server::Handle;
|
|
use std::future::Future;
|
|
use std::time::Duration;
|
|
use tokio::io::{AsyncRead, AsyncReadExt};
|
|
use tokio::time::sleep;
|
|
use tokio_util::bytes::Bytes;
|
|
|
|
pub async fn forward<'a, R, Fut, Fwd>(
|
|
session_handle: &'a Handle,
|
|
chan_id: ChannelId,
|
|
r: &mut R,
|
|
mut fwd: Fwd,
|
|
) -> Result<(), russh::Error>
|
|
where
|
|
R: AsyncRead + Send + Unpin,
|
|
Fut: Future<Output = Result<(), Bytes>> + 'a,
|
|
Fwd: FnMut(&'a Handle, ChannelId, Bytes) -> Fut,
|
|
{
|
|
const BUF_SIZE: usize = 1024 * 32;
|
|
const MAX_RETRIES: usize = 5;
|
|
const RETRY_DELAY: u64 = 10; // ms
|
|
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
loop {
|
|
let read = r.read(&mut buf).await?;
|
|
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
|
|
let mut chunk = Bytes::copy_from_slice(&buf[..read]);
|
|
let mut retries = 0;
|
|
loop {
|
|
match fwd(session_handle, chan_id, chunk).await {
|
|
Ok(()) => break,
|
|
Err(unsent) => {
|
|
retries += 1;
|
|
if retries >= MAX_RETRIES {
|
|
return Ok(());
|
|
}
|
|
chunk = unsent;
|
|
sleep(Duration::from_millis(RETRY_DELAY)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|