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