52 lines
1.3 KiB
Rust
52 lines
1.3 KiB
Rust
use std::{future::Future, time::Duration};
|
|
|
|
use actix_web::web::Bytes;
|
|
use russh::{ChannelId, server::Handle};
|
|
use tokio::{
|
|
io::{AsyncRead, AsyncReadExt},
|
|
time::sleep,
|
|
};
|
|
|
|
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;
|
|
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
loop {
|
|
let read = r.read(&mut buf).await.map_err(russh::Error::IO)?;
|
|
|
|
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(())
|
|
}
|