1use std::collections::HashMap;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14use anyhow::{Context, Result, anyhow, ensure};
15use tokio::io::{AsyncRead, AsyncWrite};
16use tokio::sync::{mpsc, oneshot};
17
18use super::{Entry, RangeReq, Refused, RemoteFs};
19use crate::sftp::transport::{self, SshChild};
20use crate::sftp::wire::{
21 Attrs, CLOSE, DATA, Dec, Enc, FXF_APPEND, FXF_CREAT, FXF_READ, FXF_WRITE, HANDLE, MKDIR, NAME,
22 OPEN, OPENDIR, READ, READDIR, STATUS, STATUS_EOF, STATUS_OK, WRITE, owner_of_longname,
23};
24use crate::sftp::{Reply, Rx, Sftp, Tx};
25
26const QUEUE_DEPTH: usize = 1024;
27const MAX_BATCH: usize = 256;
28const READ_CHUNK: u32 = 32 * 1024;
29const WRITE_CHUNK: usize = 32 * 1024;
30
31const COALESCE: Duration = Duration::from_micros(200);
36
37struct Job {
38 kind: u8,
39 body: Vec<u8>,
41 reply: oneshot::Sender<Reply>,
42}
43
44type Pending = Arc<Mutex<HashMap<u32, oneshot::Sender<Reply>>>>;
45
46pub struct SftpFs {
47 jobs: mpsc::Sender<Job>,
48 round_trips: Arc<AtomicU64>,
49 _child: Option<SshChild>,
51}
52
53impl SftpFs {
54 pub async fn connect(host: &str) -> Result<Self> {
55 let (child, w, r) = transport::open(host)?;
56 let sftp = Sftp::handshake(w, r).await?;
57 Ok(Self::drive(sftp, Some(child)))
58 }
59
60 pub async fn over<W, R>(w: W, r: R) -> Result<Self>
63 where
64 W: AsyncWrite + Unpin + Send + 'static,
65 R: AsyncRead + Unpin + Send + 'static,
66 {
67 let sftp = Sftp::handshake(w, r).await?;
68 Ok(Self::drive(sftp, None))
69 }
70
71 fn drive<W, R>(sftp: Sftp<W, R>, child: Option<SshChild>) -> Self
72 where
73 W: AsyncWrite + Unpin + Send + 'static,
74 R: AsyncRead + Unpin + Send + 'static,
75 {
76 let (tx, rx) = sftp.into_halves();
77 let (jobs, job_rx) = mpsc::channel(QUEUE_DEPTH);
78 let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
79 let round_trips = Arc::new(AtomicU64::new(0));
80
81 tokio::spawn(writer(
82 tx,
83 job_rx,
84 Arc::clone(&pending),
85 Arc::clone(&round_trips),
86 ));
87 tokio::spawn(reader(rx, pending));
88
89 Self {
90 jobs,
91 round_trips,
92 _child: child,
93 }
94 }
95
96 async fn issue(&self, kind: u8, body: Vec<u8>) -> Result<oneshot::Receiver<Reply>> {
98 let (reply, rx) = oneshot::channel();
99 self.jobs
100 .send(Job { kind, body, reply })
101 .await
102 .map_err(|_| anyhow!("sftp session is gone"))?;
103 Ok(rx)
104 }
105}
106
107async fn await_reply(rx: oneshot::Receiver<Reply>) -> Result<Reply> {
108 rx.await
109 .map_err(|_| anyhow!("sftp session closed before replying"))
110}
111
112fn decode_names(payload: &[u8]) -> Result<Vec<Entry>> {
114 let mut d = Dec::new(payload);
115 let count = d.u32().context("readdir count")?;
116 ensure!(count <= 1 << 16, "implausible readdir count {count}");
119 let mut out = Vec::with_capacity(count as usize);
120 for _ in 0..count {
121 let name = String::from_utf8_lossy(d.str().context("filename")?).into_owned();
122 let longname = String::from_utf8_lossy(d.str().context("longname")?).into_owned();
126 let owner = owner_of_longname(&longname).map(str::to_string);
127 let attrs = Attrs::decode(&mut d).context("attrs")?;
128 out.push(Entry { name, attrs, owner });
129 }
130 Ok(out)
131}
132
133fn handle_from(r: &Reply, what: &str) -> Result<Vec<u8>> {
141 if r.kind != HANDLE {
142 let why = match Dec::new(r.payload()).u32() {
143 Some(status) => anyhow::Error::new(Refused { status }),
144 None => anyhow!("unreadable reply (type {})", r.kind),
147 };
148 return Err(why.context(format!("{what} refused")));
149 }
150 Ok(Dec::new(r.payload()).str().context("handle")?.to_vec())
151}
152
153impl RemoteFs for SftpFs {
154 async fn read_batch(&self, paths: &[String]) -> Vec<Result<Vec<u8>>> {
155 let mut opens = Vec::with_capacity(paths.len());
159 for p in paths {
160 opens.push(
161 self.issue(
162 OPEN,
163 Enc::new().str(p.as_bytes()).u32(FXF_READ).u32(0).done(),
164 )
165 .await,
166 );
167 }
168
169 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(paths.len());
170 let mut out: Vec<Result<Vec<u8>>> = Vec::with_capacity(paths.len());
171 for (rx, path) in opens.into_iter().zip(paths) {
172 let opened = match rx {
173 Ok(rx) => await_reply(rx).await.and_then(|r| handle_from(&r, path)),
174 Err(e) => Err(e),
175 };
176 match opened {
177 Ok(h) => {
178 handles.push(Some(h));
179 out.push(Ok(Vec::new()));
180 }
181 Err(e) => {
182 handles.push(None);
183 out.push(Err(e));
184 }
185 }
186 }
187
188 let mut live: Vec<usize> = (0..paths.len()).filter(|&i| handles[i].is_some()).collect();
191 while !live.is_empty() {
192 let mut rxs = Vec::with_capacity(live.len());
193 for &i in &live {
194 let handle = handles[i].as_ref().expect("live implies a handle");
195 let offset = out[i].as_ref().map_or(0, Vec::len) as u64;
196 rxs.push(
197 self.issue(
198 READ,
199 Enc::new().str(handle).u64(offset).u32(READ_CHUNK).done(),
200 )
201 .await,
202 );
203 }
204
205 let mut still_live = Vec::new();
206 for (&i, rx) in live.iter().zip(rxs) {
207 let chunk = match rx {
208 Ok(rx) => await_reply(rx).await,
209 Err(e) => Err(e),
210 };
211 match chunk {
212 Ok(r) if r.kind == DATA => {
213 let data = Dec::new(r.payload()).str().unwrap_or(&[]).to_vec();
214 let full = data.len() as u32 == READ_CHUNK;
215 if let Ok(buf) = &mut out[i] {
216 buf.extend_from_slice(&data);
217 }
218 if full {
219 still_live.push(i);
220 }
221 }
222 Ok(r) if r.kind == STATUS => {
227 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
228 if code != STATUS_EOF {
229 out[i] = Err(anyhow!("read failed with sftp status {code}"));
230 }
231 }
232 Ok(r) => out[i] = Err(anyhow!("read gave reply type {}", r.kind)),
233 Err(e) => out[i] = Err(e),
234 }
235 }
236 live = still_live;
237 }
238
239 for handle in handles.iter().flatten() {
240 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
241 }
242 out
243 }
244
245 async fn read_ranges(&self, reqs: &[RangeReq]) -> Vec<Result<Vec<u8>>> {
246 let mut opens = Vec::with_capacity(reqs.len());
247 for r in reqs {
248 opens.push(
249 self.issue(
250 OPEN,
251 Enc::new()
252 .str(r.path.as_bytes())
253 .u32(FXF_READ)
254 .u32(0)
255 .done(),
256 )
257 .await,
258 );
259 }
260
261 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(reqs.len());
262 let mut out: Vec<Result<Vec<u8>>> = Vec::with_capacity(reqs.len());
263 for (rx, r) in opens.into_iter().zip(reqs) {
264 let opened = match rx {
265 Ok(rx) => await_reply(rx)
266 .await
267 .and_then(|reply| handle_from(&reply, &r.path)),
268 Err(e) => Err(e),
269 };
270 match opened {
271 Ok(h) => {
272 handles.push(Some(h));
273 out.push(Ok(Vec::new()));
274 }
275 Err(e) => {
276 handles.push(None);
277 out.push(Err(e));
278 }
279 }
280 }
281
282 struct Piece {
286 req: usize,
287 offset: u64,
288 len: u32,
289 }
290 let mut pieces = Vec::new();
291 for (i, r) in reqs.iter().enumerate() {
292 if handles[i].is_none() {
293 continue;
294 }
295 let mut at = r.offset;
296 let end = r.offset.saturating_add(r.len);
297 while at < end {
298 let len =
299 u32::try_from((end - at).min(u64::from(READ_CHUNK))).unwrap_or(READ_CHUNK);
300 pieces.push(Piece {
301 req: i,
302 offset: at,
303 len,
304 });
305 at += u64::from(len);
306 }
307 }
308
309 let mut rxs = Vec::with_capacity(pieces.len());
310 for p in &pieces {
311 let handle = handles[p.req].as_ref().expect("pieces skip failed opens");
312 rxs.push(
313 self.issue(READ, Enc::new().str(handle).u64(p.offset).u32(p.len).done())
314 .await,
315 );
316 }
317
318 for (p, rx) in pieces.iter().zip(rxs) {
321 let reply = match rx {
322 Ok(rx) => await_reply(rx).await,
323 Err(e) => Err(e),
324 };
325 match reply {
326 Ok(r) if r.kind == DATA => {
327 let data = Dec::new(r.payload()).str().unwrap_or(&[]).to_vec();
328 if let Ok(buf) = &mut out[p.req] {
329 buf.extend_from_slice(&data);
330 }
331 }
332 Ok(r) if r.kind == STATUS => {
336 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
337 if code != STATUS_EOF {
338 out[p.req] = Err(anyhow!("read failed with sftp status {code}"));
339 }
340 }
341 Ok(r) => out[p.req] = Err(anyhow!("read gave reply type {}", r.kind)),
342 Err(e) => out[p.req] = Err(e),
343 }
344 }
345
346 for handle in handles.iter().flatten() {
347 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
348 }
349 out
350 }
351
352 async fn append(&self, path: &str, bytes: &[u8]) -> Result<()> {
353 let opened = await_reply(
358 self.issue(
359 OPEN,
360 Enc::new()
361 .str(path.as_bytes())
362 .u32(FXF_WRITE | FXF_APPEND | FXF_CREAT)
363 .u32(0)
364 .done(),
365 )
366 .await?,
367 )
368 .await?;
369 let handle = handle_from(&opened, path)?;
370
371 let mut rxs = Vec::new();
373 let mut at = 0usize;
374 while at < bytes.len() {
375 let end = at.saturating_add(WRITE_CHUNK).min(bytes.len());
376 rxs.push(
377 self.issue(
378 WRITE,
379 Enc::new()
380 .str(&handle)
381 .u64(at as u64)
382 .str(&bytes[at..end])
383 .done(),
384 )
385 .await?,
386 );
387 at = end;
388 }
389
390 let mut failure = None;
394 for rx in rxs {
395 match await_reply(rx).await {
396 Ok(r) => {
397 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
398 if r.kind != STATUS || code != STATUS_OK {
399 failure = Some(anyhow!("writing {path} failed with sftp status {code}"));
400 }
401 }
402 Err(e) => failure = Some(e),
403 }
404 }
405
406 let _ = self.issue(CLOSE, Enc::new().str(&handle).done()).await;
407 match failure {
408 Some(e) => Err(e),
409 None => Ok(()),
410 }
411 }
412
413 async fn mkdirs(&self, path: &str) -> Result<()> {
414 let mut levels = Vec::new();
415 let mut at = String::new();
416 for part in path.split('/').filter(|p| !p.is_empty()) {
417 at.push('/');
418 at.push_str(part);
419 levels.push(at.clone());
420 }
421
422 let mut rxs = Vec::with_capacity(levels.len());
423 for level in &levels {
424 rxs.push(
425 self.issue(MKDIR, Enc::new().str(level.as_bytes()).u32(0).done())
426 .await?,
427 );
428 }
429 for rx in rxs {
430 let _ = await_reply(rx).await;
433 }
434
435 let one = [path.to_string()];
439 let mut got = self.list_dirs(&one).await;
440 match got.pop() {
441 Some(Ok(_)) => Ok(()),
442 Some(Err(e)) => Err(e.context(format!("creating {path}"))),
443 None => Err(anyhow!("list_dirs returned nothing for {path}")),
444 }
445 }
446
447 async fn list_dirs(&self, paths: &[String]) -> Vec<Result<Vec<Entry>>> {
448 let mut opens = Vec::with_capacity(paths.len());
453 for p in paths {
454 opens.push(
455 self.issue(OPENDIR, Enc::new().str(p.as_bytes()).done())
456 .await,
457 );
458 }
459
460 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(paths.len());
461 let mut out: Vec<Result<Vec<Entry>>> = Vec::with_capacity(paths.len());
462 for (rx, path) in opens.into_iter().zip(paths) {
463 let opened = match rx {
464 Ok(rx) => await_reply(rx).await.and_then(|r| handle_from(&r, path)),
465 Err(e) => Err(e),
466 };
467 match opened {
468 Ok(h) => {
469 handles.push(Some(h));
470 out.push(Ok(Vec::new()));
471 }
472 Err(e) => {
473 handles.push(None);
474 out.push(Err(e));
475 }
476 }
477 }
478
479 let mut live: Vec<usize> = (0..paths.len()).filter(|&i| handles[i].is_some()).collect();
483 while !live.is_empty() {
484 let mut rxs = Vec::with_capacity(live.len());
485 for &i in &live {
486 let handle = handles[i].as_ref().expect("live implies a handle");
487 rxs.push(self.issue(READDIR, Enc::new().str(handle).done()).await);
488 }
489
490 let mut still_live = Vec::new();
491 for (&i, rx) in live.iter().zip(rxs) {
492 let page = match rx {
493 Ok(rx) => await_reply(rx).await,
494 Err(e) => Err(e),
495 };
496 match page {
497 Ok(r) if r.kind == NAME => match decode_names(r.payload()) {
498 Ok(entries) => {
499 if let Ok(acc) = &mut out[i] {
500 acc.extend(entries);
501 }
502 still_live.push(i);
503 }
504 Err(e) => out[i] = Err(e),
505 },
506 Ok(r) if r.kind == STATUS => {
511 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
512 if code != STATUS_EOF {
513 out[i] = Err(anyhow!("readdir failed with sftp status {code}"));
514 }
515 }
516 Ok(r) => out[i] = Err(anyhow!("readdir gave reply type {}", r.kind)),
517 Err(e) => out[i] = Err(e),
518 }
519 }
520 live = still_live;
521 }
522
523 for handle in handles.iter().flatten() {
524 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
525 }
526 out
527 }
528
529 fn round_trips(&self) -> u64 {
530 self.round_trips.load(Ordering::Relaxed)
531 }
532}
533
534async fn writer<W: AsyncWrite + Unpin>(
535 mut tx: Tx<W>,
536 mut jobs: mpsc::Receiver<Job>,
537 pending: Pending,
538 round_trips: Arc<AtomicU64>,
539) {
540 let mut batch: Vec<Job> = Vec::with_capacity(MAX_BATCH);
541 loop {
542 if jobs.recv_many(&mut batch, MAX_BATCH).await == 0 {
543 return;
544 }
545 if batch.len() < MAX_BATCH {
546 tokio::time::sleep(COALESCE).await;
547 while batch.len() < MAX_BATCH {
548 match jobs.try_recv() {
549 Ok(job) => batch.push(job),
550 Err(_) => break,
551 }
552 }
553 }
554
555 for job in batch.drain(..) {
556 let id = tx.alloc_id();
557 let mut payload = Vec::with_capacity(4 + job.body.len());
558 payload.extend_from_slice(&id.to_be_bytes());
559 payload.extend_from_slice(&job.body);
560 pending
563 .lock()
564 .expect("pending map poisoned")
565 .insert(id, job.reply);
566 if tx.queue(job.kind, &payload).await.is_err() {
567 return;
568 }
569 }
570
571 if tx.flush().await.is_err() {
572 return;
573 }
574 round_trips.fetch_add(1, Ordering::Relaxed);
575 }
576}
577
578async fn reader<R: AsyncRead + Unpin>(mut rx: Rx<R>, pending: Pending) {
579 while let Ok(reply) = rx.recv().await {
580 let waiter = pending
581 .lock()
582 .expect("pending map poisoned")
583 .remove(&reply.id);
584 if let Some(waiter) = waiter {
585 let _ = waiter.send(reply);
586 }
587 }
588 pending.lock().expect("pending map poisoned").clear();
591}
592
593#[cfg(test)]
594mod tests {
595 use super::*;
596 use crate::sftp::wire::{INIT, VERSION};
597 use crate::sftp::{read_frame, write_frame};
598 use tokio::io::AsyncWriteExt;
599
600 async fn fake_server<R, W>(mut r: R, mut w: W, body: Vec<u8>, fail_read: bool)
603 where
604 R: AsyncRead + Unpin,
605 W: AsyncWrite + Unpin,
606 {
607 let (kind, _) = read_frame(&mut r).await.expect("init frame");
608 assert_eq!(kind, INIT);
609 write_frame(&mut w, VERSION, &Enc::new().u32(3).done())
610 .await
611 .expect("version");
612 w.flush().await.expect("flush version");
613
614 while let Ok((kind, payload)) = read_frame(&mut r).await {
615 let mut d = Dec::new(&payload);
616 let id = d.u32().expect("request id");
617 let (out_kind, out) = match kind {
618 OPEN => (HANDLE, Enc::new().u32(id).str(b"h").done()),
619 READ => {
620 d.str().expect("handle");
621 let offset = d.u64().expect("offset") as usize;
622 if fail_read {
623 (
625 STATUS,
626 Enc::new()
627 .u32(id)
628 .u32(4)
629 .str(b"is a directory")
630 .str(b"")
631 .done(),
632 )
633 } else if offset >= body.len() {
634 (
635 STATUS,
636 Enc::new().u32(id).u32(1).str(b"eof").str(b"").done(),
637 )
638 } else {
639 (DATA, Enc::new().u32(id).str(&body[offset..]).done())
640 }
641 }
642 CLOSE => (STATUS, Enc::new().u32(id).u32(0).str(b"ok").str(b"").done()),
643 other => panic!("fake server got unexpected request type {other}"),
644 };
645 write_frame(&mut w, out_kind, &out).await.expect("reply");
646 w.flush().await.expect("flush reply");
647 }
648 }
649
650 #[tokio::test]
654 async fn forty_reads_cost_a_constant_number_of_round_trips() {
655 let (client, server) = tokio::io::duplex(1 << 20);
656 let (cr, cw) = tokio::io::split(client);
657 let (sr, sw) = tokio::io::split(server);
658 tokio::spawn(fake_server(sr, sw, b"hello".to_vec(), false));
659
660 let fs = SftpFs::over(cw, cr).await.expect("handshake");
661 let paths: Vec<String> = (0..40).map(|i| format!("/f{i}")).collect();
662 let out = tokio::time::timeout(Duration::from_secs(10), fs.read_batch(&paths))
663 .await
664 .expect("read_batch should not hang");
665
666 assert_eq!(out.len(), 40);
667 for r in &out {
668 assert_eq!(r.as_ref().expect("read succeeded").as_slice(), b"hello");
669 }
670
671 let trips = fs.round_trips();
674 assert!(trips <= 6, "forty files cost {trips} round trips");
675 }
676
677 #[tokio::test]
681 async fn a_failed_read_is_not_reported_as_an_empty_success() {
682 let (client, server) = tokio::io::duplex(1 << 16);
683 let (cr, cw) = tokio::io::split(client);
684 let (sr, sw) = tokio::io::split(server);
685 tokio::spawn(fake_server(sr, sw, b"unused".to_vec(), true));
686
687 let fs = SftpFs::over(cw, cr).await.expect("handshake");
688 let out = tokio::time::timeout(
689 Duration::from_secs(10),
690 fs.read_batch(&["/a-directory".to_string()]),
691 )
692 .await
693 .expect("read_batch should not hang");
694
695 assert!(
696 out[0].is_err(),
697 "a read that failed must not look like an empty file"
698 );
699 }
700
701 #[tokio::test]
704 async fn a_dead_session_fails_callers_instead_of_hanging() {
705 let (client, server) = tokio::io::duplex(1 << 16);
706 let (cr, cw) = tokio::io::split(client);
707 let (mut sr, mut sw) = tokio::io::split(server);
708 tokio::spawn(async move {
709 read_frame(&mut sr).await.expect("init frame");
710 write_frame(&mut sw, VERSION, &Enc::new().u32(3).done())
711 .await
712 .expect("version");
713 sw.flush().await.expect("flush version");
714 });
716
717 let fs = SftpFs::over(cw, cr).await.expect("handshake");
718 let out = tokio::time::timeout(Duration::from_secs(10), fs.read_batch(&["/a".to_string()]))
719 .await
720 .expect("a dead session must not hang the caller");
721
722 assert!(out[0].is_err(), "a closed session must surface as an error");
723 }
724}