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_READ, HANDLE, NAME, OPEN, OPENDIR, READ, READDIR, REALPATH,
22 STATUS, STATUS_EOF, Verb,
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;
29
30const COALESCE: Duration = Duration::from_micros(200);
35
36struct Job {
37 kind: Verb,
38 body: Vec<u8>,
40 reply: oneshot::Sender<Reply>,
41}
42
43type Pending = Arc<Mutex<HashMap<u32, oneshot::Sender<Reply>>>>;
44
45pub struct SftpFs {
46 jobs: mpsc::Sender<Job>,
47 round_trips: Arc<AtomicU64>,
48 _child: Option<SshChild>,
50}
51
52impl SftpFs {
53 pub async fn connect(host: &str) -> Result<Self> {
54 let (child, w, r) = transport::open(host)?;
55 let sftp = Sftp::handshake(w, r).await?;
56 Ok(Self::drive(sftp, Some(child)))
57 }
58
59 pub fn is_alive(&self) -> bool {
70 !self.jobs.is_closed()
71 }
72
73 pub async fn over<W, R>(w: W, r: R) -> Result<Self>
76 where
77 W: AsyncWrite + Unpin + Send + 'static,
78 R: AsyncRead + Unpin + Send + 'static,
79 {
80 let sftp = Sftp::handshake(w, r).await?;
81 Ok(Self::drive(sftp, None))
82 }
83
84 fn drive<W, R>(sftp: Sftp<W, R>, child: Option<SshChild>) -> Self
85 where
86 W: AsyncWrite + Unpin + Send + 'static,
87 R: AsyncRead + Unpin + Send + 'static,
88 {
89 let (tx, rx) = sftp.into_halves();
90 let (jobs, job_rx) = mpsc::channel(QUEUE_DEPTH);
91 let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
92 let round_trips = Arc::new(AtomicU64::new(0));
93
94 tokio::spawn(writer(
95 tx,
96 job_rx,
97 Arc::clone(&pending),
98 Arc::clone(&round_trips),
99 ));
100 tokio::spawn(reader(rx, pending));
101
102 Self {
103 jobs,
104 round_trips,
105 _child: child,
106 }
107 }
108
109 async fn issue(&self, kind: Verb, body: Vec<u8>) -> Result<oneshot::Receiver<Reply>> {
111 let (reply, rx) = oneshot::channel();
112 self.jobs
113 .send(Job { kind, body, reply })
114 .await
115 .map_err(|_| anyhow!("sftp session is gone"))?;
116 Ok(rx)
117 }
118}
119
120async fn await_reply(rx: oneshot::Receiver<Reply>) -> Result<Reply> {
121 rx.await
122 .map_err(|_| anyhow!("sftp session closed before replying"))
123}
124
125fn decode_names(payload: &[u8]) -> Result<Vec<Entry>> {
127 let mut d = Dec::new(payload);
128 let count = d.u32().context("readdir count")?;
129 ensure!(count <= 1 << 16, "implausible readdir count {count}");
132 let mut out = Vec::with_capacity(count as usize);
133 for _ in 0..count {
134 let name = String::from_utf8_lossy(d.str().context("filename")?).into_owned();
135 d.str().context("longname")?;
138 let attrs = Attrs::decode(&mut d).context("attrs")?;
139 out.push(Entry { name, attrs });
140 }
141 Ok(out)
142}
143
144fn handle_from(r: &Reply, what: &str) -> Result<Vec<u8>> {
151 if r.kind != HANDLE {
152 let why = match Dec::new(r.payload()).u32() {
153 Some(status) => anyhow::Error::new(Refused { status }),
154 None => anyhow!("unreadable reply (type {})", r.kind),
157 };
158 return Err(why.context(format!("{what} refused")));
159 }
160 Ok(Dec::new(r.payload()).str().context("handle")?.to_vec())
161}
162
163impl RemoteFs for SftpFs {
164 async fn read_batch(&self, paths: &[String]) -> Vec<Result<Vec<u8>>> {
165 let mut opens = Vec::with_capacity(paths.len());
169 for p in paths {
170 opens.push(
171 self.issue(
172 OPEN,
173 Enc::new().str(p.as_bytes()).u32(FXF_READ).u32(0).done(),
174 )
175 .await,
176 );
177 }
178
179 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(paths.len());
180 let mut out: Vec<Result<Vec<u8>>> = Vec::with_capacity(paths.len());
181 for (rx, path) in opens.into_iter().zip(paths) {
182 let opened = match rx {
183 Ok(rx) => await_reply(rx).await.and_then(|r| handle_from(&r, path)),
184 Err(e) => Err(e),
185 };
186 match opened {
187 Ok(h) => {
188 handles.push(Some(h));
189 out.push(Ok(Vec::new()));
190 }
191 Err(e) => {
192 handles.push(None);
193 out.push(Err(e));
194 }
195 }
196 }
197
198 let mut live: Vec<usize> = (0..paths.len()).filter(|&i| handles[i].is_some()).collect();
201 while !live.is_empty() {
202 let mut rxs = Vec::with_capacity(live.len());
203 for &i in &live {
204 let handle = handles[i].as_ref().expect("live implies a handle");
205 let offset = out[i].as_ref().map_or(0, Vec::len) as u64;
206 rxs.push(
207 self.issue(
208 READ,
209 Enc::new().str(handle).u64(offset).u32(READ_CHUNK).done(),
210 )
211 .await,
212 );
213 }
214
215 let mut still_live = Vec::new();
216 for (&i, rx) in live.iter().zip(rxs) {
217 let chunk = match rx {
218 Ok(rx) => await_reply(rx).await,
219 Err(e) => Err(e),
220 };
221 match chunk {
222 Ok(r) if r.kind == DATA => {
223 let data = Dec::new(r.payload()).str().unwrap_or(&[]).to_vec();
224 let full = data.len() as u32 == READ_CHUNK;
225 if let Ok(buf) = &mut out[i] {
226 buf.extend_from_slice(&data);
227 }
228 if full {
229 still_live.push(i);
230 }
231 }
232 Ok(r) if r.kind == STATUS => {
237 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
238 if code != STATUS_EOF {
239 out[i] = Err(anyhow!("read failed with sftp status {code}"));
240 }
241 }
242 Ok(r) => out[i] = Err(anyhow!("read gave reply type {}", r.kind)),
243 Err(e) => out[i] = Err(e),
244 }
245 }
246 live = still_live;
247 }
248
249 for handle in handles.iter().flatten() {
250 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
251 }
252 out
253 }
254
255 async fn read_ranges(&self, reqs: &[RangeReq]) -> Vec<Result<Vec<u8>>> {
256 let mut opens = Vec::with_capacity(reqs.len());
257 for r in reqs {
258 opens.push(
259 self.issue(
260 OPEN,
261 Enc::new()
262 .str(r.path.as_bytes())
263 .u32(FXF_READ)
264 .u32(0)
265 .done(),
266 )
267 .await,
268 );
269 }
270
271 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(reqs.len());
272 let mut out: Vec<Result<Vec<u8>>> = Vec::with_capacity(reqs.len());
273 for (rx, r) in opens.into_iter().zip(reqs) {
274 let opened = match rx {
275 Ok(rx) => await_reply(rx)
276 .await
277 .and_then(|reply| handle_from(&reply, &r.path)),
278 Err(e) => Err(e),
279 };
280 match opened {
281 Ok(h) => {
282 handles.push(Some(h));
283 out.push(Ok(Vec::new()));
284 }
285 Err(e) => {
286 handles.push(None);
287 out.push(Err(e));
288 }
289 }
290 }
291
292 struct Piece {
296 req: usize,
297 offset: u64,
298 len: u32,
299 }
300 let mut pieces = Vec::new();
301 for (i, r) in reqs.iter().enumerate() {
302 if handles[i].is_none() {
303 continue;
304 }
305 let mut at = r.offset;
306 let end = r.offset.saturating_add(r.len);
307 while at < end {
308 let len =
309 u32::try_from((end - at).min(u64::from(READ_CHUNK))).unwrap_or(READ_CHUNK);
310 pieces.push(Piece {
311 req: i,
312 offset: at,
313 len,
314 });
315 at += u64::from(len);
316 }
317 }
318
319 let mut rxs = Vec::with_capacity(pieces.len());
320 for p in &pieces {
321 let handle = handles[p.req].as_ref().expect("pieces skip failed opens");
322 rxs.push(
323 self.issue(READ, Enc::new().str(handle).u64(p.offset).u32(p.len).done())
324 .await,
325 );
326 }
327
328 for (p, rx) in pieces.iter().zip(rxs) {
331 let reply = match rx {
332 Ok(rx) => await_reply(rx).await,
333 Err(e) => Err(e),
334 };
335 match reply {
336 Ok(r) if r.kind == DATA => {
337 let data = Dec::new(r.payload()).str().unwrap_or(&[]).to_vec();
338 if let Ok(buf) = &mut out[p.req] {
339 buf.extend_from_slice(&data);
340 }
341 }
342 Ok(r) if r.kind == STATUS => {
346 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
347 if code != STATUS_EOF {
348 out[p.req] = Err(anyhow!("read failed with sftp status {code}"));
349 }
350 }
351 Ok(r) => out[p.req] = Err(anyhow!("read gave reply type {}", r.kind)),
352 Err(e) => out[p.req] = Err(e),
353 }
354 }
355
356 for handle in handles.iter().flatten() {
357 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
358 }
359 out
360 }
361
362 async fn home(&self) -> Result<String> {
363 let rx = self.issue(REALPATH, Enc::new().str(b".").done()).await?;
369 let reply = await_reply(rx).await?;
370 ensure!(
371 reply.kind == NAME,
372 "realpath answered {} rather than a name",
373 reply.kind
374 );
375 let mut d = Dec::new(reply.payload());
376 let count = d.u32().context("realpath count")?;
380 ensure!(count >= 1, "realpath answered with no name");
381 let path = String::from_utf8(d.str().context("realpath name")?.to_vec())
382 .context("home directory path is not utf-8")?;
383 ensure!(
387 path.starts_with('/'),
388 "realpath answered {path:?}, which is not an absolute path"
389 );
390 Ok(path)
391 }
392
393 async fn list_dirs(&self, paths: &[String]) -> Vec<Result<Vec<Entry>>> {
394 let mut opens = Vec::with_capacity(paths.len());
399 for p in paths {
400 opens.push(
401 self.issue(OPENDIR, Enc::new().str(p.as_bytes()).done())
402 .await,
403 );
404 }
405
406 let mut handles: Vec<Option<Vec<u8>>> = Vec::with_capacity(paths.len());
407 let mut out: Vec<Result<Vec<Entry>>> = Vec::with_capacity(paths.len());
408 for (rx, path) in opens.into_iter().zip(paths) {
409 let opened = match rx {
410 Ok(rx) => await_reply(rx).await.and_then(|r| handle_from(&r, path)),
411 Err(e) => Err(e),
412 };
413 match opened {
414 Ok(h) => {
415 handles.push(Some(h));
416 out.push(Ok(Vec::new()));
417 }
418 Err(e) => {
419 handles.push(None);
420 out.push(Err(e));
421 }
422 }
423 }
424
425 let mut live: Vec<usize> = (0..paths.len()).filter(|&i| handles[i].is_some()).collect();
429 while !live.is_empty() {
430 let mut rxs = Vec::with_capacity(live.len());
431 for &i in &live {
432 let handle = handles[i].as_ref().expect("live implies a handle");
433 rxs.push(self.issue(READDIR, Enc::new().str(handle).done()).await);
434 }
435
436 let mut still_live = Vec::new();
437 for (&i, rx) in live.iter().zip(rxs) {
438 let page = match rx {
439 Ok(rx) => await_reply(rx).await,
440 Err(e) => Err(e),
441 };
442 match page {
443 Ok(r) if r.kind == NAME => match decode_names(r.payload()) {
444 Ok(entries) => {
445 if let Ok(acc) = &mut out[i] {
446 acc.extend(entries);
447 }
448 still_live.push(i);
449 }
450 Err(e) => out[i] = Err(e),
451 },
452 Ok(r) if r.kind == STATUS => {
457 let code = Dec::new(r.payload()).u32().unwrap_or(u32::MAX);
458 if code != STATUS_EOF {
459 out[i] = Err(anyhow!("readdir failed with sftp status {code}"));
460 }
461 }
462 Ok(r) => out[i] = Err(anyhow!("readdir gave reply type {}", r.kind)),
463 Err(e) => out[i] = Err(e),
464 }
465 }
466 live = still_live;
467 }
468
469 for handle in handles.iter().flatten() {
470 let _ = self.issue(CLOSE, Enc::new().str(handle).done()).await;
471 }
472 out
473 }
474
475 fn round_trips(&self) -> u64 {
476 self.round_trips.load(Ordering::Relaxed)
477 }
478}
479
480async fn writer<W: AsyncWrite + Unpin>(
481 mut tx: Tx<W>,
482 mut jobs: mpsc::Receiver<Job>,
483 pending: Pending,
484 round_trips: Arc<AtomicU64>,
485) {
486 let mut batch: Vec<Job> = Vec::with_capacity(MAX_BATCH);
487 loop {
488 if jobs.recv_many(&mut batch, MAX_BATCH).await == 0 {
489 return;
490 }
491 if batch.len() < MAX_BATCH {
492 tokio::time::sleep(COALESCE).await;
493 while batch.len() < MAX_BATCH {
494 match jobs.try_recv() {
495 Ok(job) => batch.push(job),
496 Err(_) => break,
497 }
498 }
499 }
500
501 for job in batch.drain(..) {
502 let id = tx.alloc_id();
503 let mut payload = Vec::with_capacity(4 + job.body.len());
504 payload.extend_from_slice(&id.to_be_bytes());
505 payload.extend_from_slice(&job.body);
506 pending
509 .lock()
510 .expect("pending map poisoned")
511 .insert(id, job.reply);
512 if tx.queue(job.kind, &payload).await.is_err() {
513 return;
514 }
515 }
516
517 if tx.flush().await.is_err() {
518 return;
519 }
520 round_trips.fetch_add(1, Ordering::Relaxed);
521 }
522}
523
524async fn reader<R: AsyncRead + Unpin>(mut rx: Rx<R>, pending: Pending) {
525 while let Ok(reply) = rx.recv().await {
526 let waiter = pending
527 .lock()
528 .expect("pending map poisoned")
529 .remove(&reply.id);
530 if let Some(waiter) = waiter {
531 let _ = waiter.send(reply);
532 }
533 }
534 pending.lock().expect("pending map poisoned").clear();
537}
538
539#[cfg(test)]
540mod tests {
541 use super::*;
542 use crate::sftp::wire::{INIT, VERSION};
543
544 const B_OPEN: u8 = OPEN.code();
550 const B_CLOSE: u8 = CLOSE.code();
551 const B_READ: u8 = READ.code();
552 use crate::sftp::{read_frame, write_frame};
553 use tokio::io::AsyncWriteExt;
554
555 async fn fake_server<R, W>(mut r: R, mut w: W, body: Vec<u8>, fail_read: bool)
558 where
559 R: AsyncRead + Unpin,
560 W: AsyncWrite + Unpin,
561 {
562 let (kind, _) = read_frame(&mut r).await.expect("init frame");
563 assert_eq!(kind, INIT);
564 write_frame(&mut w, VERSION, &Enc::new().u32(3).done())
565 .await
566 .expect("version");
567 w.flush().await.expect("flush version");
568
569 while let Ok((kind, payload)) = read_frame(&mut r).await {
570 let mut d = Dec::new(&payload);
571 let id = d.u32().expect("request id");
572 let (out_kind, out) = match kind {
573 B_OPEN => (HANDLE, Enc::new().u32(id).str(b"h").done()),
574 B_READ => {
575 d.str().expect("handle");
576 let offset = d.u64().expect("offset") as usize;
577 if fail_read {
578 (
580 STATUS,
581 Enc::new()
582 .u32(id)
583 .u32(4)
584 .str(b"is a directory")
585 .str(b"")
586 .done(),
587 )
588 } else if offset >= body.len() {
589 (
590 STATUS,
591 Enc::new().u32(id).u32(1).str(b"eof").str(b"").done(),
592 )
593 } else {
594 (DATA, Enc::new().u32(id).str(&body[offset..]).done())
595 }
596 }
597 B_CLOSE => (STATUS, Enc::new().u32(id).u32(0).str(b"ok").str(b"").done()),
598 other => panic!("fake server got unexpected request type {other}"),
599 };
600 write_frame(&mut w, out_kind, &out).await.expect("reply");
601 w.flush().await.expect("flush reply");
602 }
603 }
604
605 #[tokio::test]
609 async fn forty_reads_cost_a_constant_number_of_round_trips() {
610 let (client, server) = tokio::io::duplex(1 << 20);
611 let (cr, cw) = tokio::io::split(client);
612 let (sr, sw) = tokio::io::split(server);
613 tokio::spawn(fake_server(sr, sw, b"hello".to_vec(), false));
614
615 let fs = SftpFs::over(cw, cr).await.expect("handshake");
616 let paths: Vec<String> = (0..40).map(|i| format!("/f{i}")).collect();
617 let out = tokio::time::timeout(Duration::from_secs(10), fs.read_batch(&paths))
618 .await
619 .expect("read_batch should not hang");
620
621 assert_eq!(out.len(), 40);
622 for r in &out {
623 assert_eq!(r.as_ref().expect("read succeeded").as_slice(), b"hello");
624 }
625
626 let trips = fs.round_trips();
629 assert!(trips <= 6, "forty files cost {trips} round trips");
630 }
631
632 #[tokio::test]
636 async fn a_failed_read_is_not_reported_as_an_empty_success() {
637 let (client, server) = tokio::io::duplex(1 << 16);
638 let (cr, cw) = tokio::io::split(client);
639 let (sr, sw) = tokio::io::split(server);
640 tokio::spawn(fake_server(sr, sw, b"unused".to_vec(), true));
641
642 let fs = SftpFs::over(cw, cr).await.expect("handshake");
643 let out = tokio::time::timeout(
644 Duration::from_secs(10),
645 fs.read_batch(&["/a-directory".to_string()]),
646 )
647 .await
648 .expect("read_batch should not hang");
649
650 assert!(
651 out[0].is_err(),
652 "a read that failed must not look like an empty file"
653 );
654 }
655
656 #[tokio::test]
659 async fn a_dead_session_fails_callers_instead_of_hanging() {
660 let (client, server) = tokio::io::duplex(1 << 16);
661 let (cr, cw) = tokio::io::split(client);
662 let (mut sr, mut sw) = tokio::io::split(server);
663 tokio::spawn(async move {
664 read_frame(&mut sr).await.expect("init frame");
665 write_frame(&mut sw, VERSION, &Enc::new().u32(3).done())
666 .await
667 .expect("version");
668 sw.flush().await.expect("flush version");
669 });
671
672 let fs = SftpFs::over(cw, cr).await.expect("handshake");
673 let out = tokio::time::timeout(Duration::from_secs(10), fs.read_batch(&["/a".to_string()]))
674 .await
675 .expect("a dead session must not hang the caller");
676
677 assert!(out[0].is_err(), "a closed session must surface as an error");
678 }
679
680 #[tokio::test]
686 async fn the_home_directory_is_whatever_the_remote_says_it_is() {
687 let fs = crate::testing::FakeRemote::new()
688 .home("/export/scratch/u42")
689 .spawn()
690 .await;
691 assert_eq!(fs.home().await.expect("a home"), "/export/scratch/u42");
692 }
693
694 #[tokio::test]
697 async fn a_remote_that_will_not_say_where_home_is_fails_rather_than_guessing() {
698 let fs = crate::testing::FakeRemote::new().spawn().await;
699 assert!(
700 fs.home().await.is_err(),
701 "a refusal must not turn into a default path"
702 );
703 }
704}