diff --git a/Cargo.lock b/Cargo.lock index c7aa4a0..f46e7eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -96,7 +96,7 @@ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "memona-plugin-ftp" -version = "0.1.0" +version = "0.1.1" dependencies = [ "memona-plugin-sdk", ] diff --git a/Cargo.toml b/Cargo.toml index e1ad925..3265f3f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "memona-plugin-ftp" -version = "0.1.0" +version = "0.1.1" edition = "2024" license = "MIT" diff --git a/README.md b/README.md index 9d22a85..ebd7ee2 100644 --- a/README.md +++ b/README.md @@ -14,9 +14,10 @@ One provider, `ftp`, for the `ftp` and `ftps` schemes. A connection is configured with a server name, an optional port, a TLS toggle for FTPS, and an optional username and password. An empty username signs in anonymously. -Listing needs a server that answers `MLSD`, the machine-readable listing -command. A server without it reports the folder as unsupported rather than -guessing at the older, human-formatted listing. FTP carries no version token, +Listing and stat need a server that answers `MLSD` and `MLST`, the +machine-readable listing commands. A server without them reports the folder as +unsupported rather than guessing at the older, human-formatted listing or +changing its working directory to find out. FTP carries no version token, so a conditional write is reported as unsupported rather than simulated. ## Passive mode and the port range diff --git a/manifest.json b/manifest.json index a4a1ee6..34ad259 100644 --- a/manifest.json +++ b/manifest.json @@ -2,7 +2,7 @@ "schemaVersion": 1, "id": "memona/ftp", "name": "FTP", - "version": "0.1.0", + "version": "0.1.1", "apiVersion": "0.1.0", "description": "Open an FTP server as a folder, and let agents work in it.", "license": "MIT", diff --git a/src/control.rs b/src/control.rs index d4dd3e2..0fe3f5b 100644 --- a/src/control.rs +++ b/src/control.rs @@ -22,14 +22,29 @@ pub fn outcome(code: u16) -> Result<(), ErrorCode> { 200..=399 => Ok(()), // 530 is a refused login and 550 covers both a missing file and a denied // one; the server does not distinguish them, so neither does this. - 530 | 532 => Err(ErrorCode::Denied), + 530 | 532 | 553 => Err(ErrorCode::Denied), 550 => Err(ErrorCode::NotFound), 450 | 452 | 552 => Err(ErrorCode::Conflict), - 502 | 504 => Err(ErrorCode::Unsupported), + 500 | 502 | 504 => Err(ErrorCode::Unsupported), + // An invalid REST offset, such as one past the end of the file. + 501 | 554 => Err(ErrorCode::InvalidInput), _ => Err(ErrorCode::Unavailable), } } +/// What the first reply to a transfer command means. A transfer answers with +/// a preliminary reply while the data connection is open; its completion +/// comes only after that connection closes and is read separately. +pub fn preliminary(code: u16) -> Result<(), ErrorCode> { + match code { + 100..=199 => Ok(()), + // A transfer that skips the preliminary reply leaves nothing to read once + // the data connection closes, so it is refused rather than guessed at. + 200..=399 => Err(ErrorCode::Unavailable), + other => outcome(other), + } +} + pub struct Control { socket: host::Tcp, host: String, @@ -76,21 +91,25 @@ impl Control { } } - /// Reads one complete reply, following the `NNN-` continuation form. + /// Reads one complete reply, following the `NNN-` continuation form. The + /// text keeps every line, since an MLST reply carries its facts on one. pub fn reply(&mut self) -> Result<(u16, String), ErrorCode> { - let first = self.line()?; - let code: u16 = first + let mut text = self.line()?; + let code: u16 = text .get(..3) - .and_then(|text| text.parse().ok()) + .and_then(|digits| digits.parse().ok()) .ok_or(ErrorCode::Unavailable)?; - if first.as_bytes().get(3) != Some(&b'-') { - return Ok((code, first)); + if text.as_bytes().get(3) != Some(&b'-') { + return Ok((code, text)); } let terminator = format!("{code} "); for _ in 0..MAX_REPLY_LINES { let line = self.line()?; - if line.starts_with(&terminator) { - return Ok((code, line)); + let last = line.starts_with(&terminator); + text.push('\n'); + text.push_str(&line); + if last { + return Ok((code, text)); } } Err(ErrorCode::Unavailable) @@ -130,6 +149,16 @@ impl Control { let port = extended_port(&text).ok_or(ErrorCode::Unavailable)?; host::connect(&self.host, port, false).map_err(transport) } + + /// Ends a transfer whose completion reply is still owed, once its data + /// connection is closed. RFC 959 has the server answer with one reply for + /// the transfer and one for the ABOR itself, whether it had finished or + /// not, so both are read and the channel is in sync for the next command. + pub fn abort(&mut self) -> Result<(), ErrorCode> { + self.send("ABOR")?; + self.reply()?; + Ok(()) + } } /// Reads the port out of `229 Entering Extended Passive Mode (|||1234|)`. @@ -147,7 +176,7 @@ fn extended_port(text: &str) -> Option { #[cfg(test)] mod tests { - use super::{extended_port, outcome}; + use super::{extended_port, outcome, preliminary}; use memona_plugin_sdk::ErrorCode; #[test] @@ -167,9 +196,24 @@ mod tests { fn reply_codes_keep_refused_and_missing_apart() { assert!(outcome(226).is_ok()); assert!(matches!(outcome(530), Err(ErrorCode::Denied))); + assert!(matches!(outcome(553), Err(ErrorCode::Denied))); assert!(matches!(outcome(550), Err(ErrorCode::NotFound))); assert!(matches!(outcome(452), Err(ErrorCode::Conflict))); assert!(matches!(outcome(502), Err(ErrorCode::Unsupported))); + assert!(matches!(outcome(554), Err(ErrorCode::InvalidInput))); assert!(matches!(outcome(421), Err(ErrorCode::Unavailable))); + // A preliminary reply is not a completion: the plain classification + // treats it as a failure, so it must not be used on transfer commands. + assert!(matches!(outcome(150), Err(ErrorCode::Unavailable))); + } + + #[test] + fn a_transfer_starts_on_a_preliminary_reply_only() { + assert!(preliminary(125).is_ok()); + assert!(preliminary(150).is_ok()); + assert!(matches!(preliminary(226), Err(ErrorCode::Unavailable))); + assert!(matches!(preliminary(550), Err(ErrorCode::NotFound))); + assert!(matches!(preliminary(553), Err(ErrorCode::Denied))); + assert!(matches!(preliminary(425), Err(ErrorCode::Unavailable))); } } diff --git a/src/lib.rs b/src/lib.rs index 4fdffb4..be89908 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,22 +2,29 @@ mod control; mod listing; +mod transfer; -use control::{Control, outcome, transport}; -use listing::machine_entry; +use control::{Control, outcome, preliminary, transport}; +use listing::{machine_info, page}; use memona_plugin_sdk::memona::filesystem::host; use memona_plugin_sdk::{ - Capabilities, ChangePage, Entry, EntryKind, ErrorCode, FileInfo, Guest, - GuestReader, GuestWriter, ListPage, Locator, Reader, Setting, WriteCondition, - Writer, + Capabilities, ChangePage, EntryKind, ErrorCode, FileInfo, Guest, GuestReader, + GuestWriter, ListPage, Locator, Reader, Setting, WriteCondition, Writer, }; -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; +use std::collections::HashMap; +use std::rc::Rc; +use transfer::Buffered; const MAX_CHUNK: u32 = 1024 * 1024; -const MAX_ENTRIES: usize = 4096; +/// One listing may be this large, and so may every listing kept together. +const MAX_LISTING_BYTES: usize = 8 * 1024 * 1024; thread_local! { static CONTROL: RefCell> = const { RefCell::new(None) }; + /// Raw MLSD listings by path, so later pages of a folder are parsed from + /// memory instead of downloaded and parsed whole again. + static LISTINGS: RefCell>> = RefCell::new(HashMap::new()); } struct Ftp; @@ -48,18 +55,19 @@ fn with_control( }) } -/// Starts a passive transfer and leaves the data socket open for the caller. +/// Starts a passive transfer and leaves the data socket open for the caller, +/// who owes the server one more reply once that socket is closed: either the +/// completion, read by `settle`, or the answer to an `abort`. fn transfer(command: &str) -> Result { with_control(|control| { let data = control.passive()?; let (code, _) = control.send(command)?; - outcome(code)?; + preliminary(code)?; Ok(data) }) } -/// Consumes the reply that follows a completed transfer, so the next command -/// does not read someone else's answer. +/// Consumes the completion reply of a transfer whose data socket is closed. fn settle() -> Result<(), ErrorCode> { with_control(|control| { let (code, _) = control.reply()?; @@ -67,6 +75,18 @@ fn settle() -> Result<(), ErrorCode> { }) } +/// Ends a transfer early, after its data socket is closed. +fn abort() -> Result<(), ErrorCode> { + with_control(Control::abort) +} + +/// Names the kind and size of a path without moving the server's working +/// directory, which a CWD probe would do and never undo. +fn probe(control: &mut Control, path: &str) -> Result { + let reply = control.expect(&format!("MLST {path}"))?; + machine_info(&reply).ok_or(ErrorCode::Unavailable) +} + /// Reads a data socket to the end, bounded so one listing cannot grow forever. fn drain(socket: &host::Tcp, limit: usize) -> Result, ErrorCode> { let mut data = Vec::new(); @@ -82,20 +102,49 @@ fn drain(socket: &host::Tcp, limit: usize) -> Result, ErrorCode> { } } -/// Reads one file from a passive data socket. +/// Downloads a folder's MLSD listing and keeps it for its later pages. +fn fetch_listing(path: &str) -> Result, ErrorCode> { + let socket = transfer(&format!("MLSD {path}"))?; + let data = match drain(&socket, MAX_LISTING_BYTES) { + Ok(data) => data, + Err(error) => { + drop(socket); + let _ = abort(); + return Err(error); + } + }; + drop(socket); + settle()?; + let listing: Rc<[u8]> = data.into(); + LISTINGS.with_borrow_mut(|listings| { + let held: usize = listings.values().map(|kept| kept.len()).sum(); + if held + listing.len() > MAX_LISTING_BYTES { + listings.clear(); + } + listings.insert(path.to_owned(), listing.clone()); + }); + Ok(listing) +} + +/// Reads one file from a passive data socket. The server's completion reply +/// is owed from the moment the transfer starts until the socket closes, and +/// this reader consumes it itself, at the end or when dropped early. struct FtpReader { - socket: host::Tcp, - settled: RefCell, + socket: RefCell>, + buffered: RefCell, } impl GuestReader for FtpReader { fn pull(&self, limit: u32) -> Result, ErrorCode> { - if *self.settled.borrow() { - return Ok(Vec::new()); - } - let data = self.socket.read(limit.min(MAX_CHUNK)).map_err(transport)?; - if data.is_empty() { - *self.settled.borrow_mut() = true; + let mut buffered = self.buffered.borrow_mut(); + let data = buffered.pull(limit as usize, |wanted| { + let wanted = u32::try_from(wanted).unwrap_or(MAX_CHUNK).min(MAX_CHUNK); + match self.socket.borrow().as_ref() { + Some(socket) => socket.read(wanted).map_err(transport), + None => Ok(Vec::new()), + } + })?; + if buffered.ended() && self.socket.take().is_some() { settle()?; } Ok(data) @@ -106,44 +155,67 @@ impl GuestReader for FtpReader { } } +impl Drop for FtpReader { + /// A reader dropped before the end still owes the server its reply. + fn drop(&mut self) { + if let Some(socket) = self.socket.take() { + drop(socket); + let _ = abort(); + } + } +} + /// Writes one file to a passive data socket. struct FtpWriter { - socket: host::Tcp, - done: RefCell, + path: String, + socket: RefCell>, + written: Cell, } impl GuestWriter for FtpWriter { fn push(&self, data: Vec) -> Result<(), ErrorCode> { - if *self.done.borrow() { - return Err(ErrorCode::InvalidInput); - } + let socket = self.socket.borrow(); + let socket = socket.as_ref().ok_or(ErrorCode::InvalidInput)?; let mut sent = 0; while sent < data.len() { - let written = - self.socket.write(&data[sent..]).map_err(transport)? as usize; + let written = socket.write(&data[sent..]).map_err(transport)? as usize; if written == 0 { return Err(ErrorCode::Unavailable); } sent += written; } + self.written.set(self.written.get() + data.len() as u64); Ok(()) } - /// Closing the data socket is what ends an FTP upload, so the handle is - /// dropped here and the control reply read once. + /// Closing the data socket is what ends an FTP upload; only then does the + /// server answer, so the socket is given up before the reply is read. fn finish(&self) -> Result, ErrorCode> { - if *self.done.borrow() { - return Err(ErrorCode::InvalidInput); - } - *self.done.borrow_mut() = true; + drop(self.socket.take().ok_or(ErrorCode::InvalidInput)?); settle()?; Ok(Some(FileInfo { kind: EntryKind::File, - size: 0, + size: self.written.get(), revision: None, })) } } +impl Drop for FtpWriter { + /// An upload abandoned before it finished must leave no partial file. The + /// server takes the closed socket as the end of the upload, so what it + /// stored is deleted once the channel is back in sync. + fn drop(&mut self) { + let Some(socket) = self.socket.take() else { + return; + }; + drop(socket); + if abort().is_ok() { + let _ = + with_control(|control| control.expect(&format!("DELE {}", self.path))); + } + } +} + impl Guest for Ftp { type Reader = FtpReader; type Writer = FtpWriter; @@ -203,8 +275,6 @@ impl Guest for Ftp { }); } - /// FTP reports a size and a directory separately, so a stat is SIZE first - /// and a directory change only when that fails. fn stat(target: Locator) -> Result { let path = route(&target)?; if path == "/" { @@ -214,57 +284,30 @@ impl Guest for Ftp { revision: None, }); } - with_control(|control| { - let (code, text) = control.send(&format!("SIZE {path}"))?; - if (200..300).contains(&code) { - let size = text - .split_whitespace() - .nth(1) - .and_then(|value| value.parse().ok()) - .unwrap_or(0); - return Ok(FileInfo { - kind: EntryKind::File, - size, - revision: None, - }); - } - let (code, _) = control.send(&format!("CWD {path}"))?; - outcome(code)?; - Ok(FileInfo { - kind: EntryKind::Directory, - size: 0, - revision: None, - }) - }) + with_control(|control| probe(control, &path)) } /// MLSD is the machine-readable listing; a server without it cannot be read /// reliably, so this reports unsupported rather than guessing at LIST text. + /// The first page downloads the listing; later pages read the kept copy. fn list( target: Locator, cursor: Option, limit: u32, ) -> Result { let path = route(&target)?; - let start: usize = match cursor { + let start: usize = match &cursor { Some(value) => value.parse().map_err(|_| ErrorCode::InvalidInput)?, None => 0, }; - let socket = transfer(&format!("MLSD {path}"))?; - let data = drain(&socket, 8 * 1024 * 1024)?; - drop(socket); - settle()?; - let children: Vec = String::from_utf8_lossy(&data) - .lines() - .filter_map(machine_entry) - .take(MAX_ENTRIES) - .collect(); - let end = start.saturating_add(limit as usize).min(children.len()); - let entries = children.get(start..end).unwrap_or_default().to_vec(); - Ok(ListPage { - entries, - cursor: (end < children.len()).then(|| end.to_string()), - }) + let kept = cursor.as_ref().and_then(|_| { + LISTINGS.with_borrow(|listings| listings.get(&path).cloned()) + }); + let listing = match kept { + Some(listing) => listing, + None => fetch_listing(&path)?, + }; + page(&listing, start, limit as usize) } fn open_read(target: Locator, offset: u64) -> Result { @@ -274,8 +317,8 @@ impl Guest for Ftp { } let socket = transfer(&format!("RETR {path}"))?; Ok(Reader::new(FtpReader { - socket, - settled: RefCell::new(false), + socket: RefCell::new(Some(socket)), + buffered: RefCell::new(Buffered::default()), })) } @@ -290,8 +333,9 @@ impl Guest for Ftp { } let socket = transfer(&format!("STOR {path}"))?; Ok(Writer::new(FtpWriter { - socket, - done: RefCell::new(false), + path, + socket: RefCell::new(Some(socket)), + written: Cell::new(0), })) } @@ -305,16 +349,23 @@ impl Guest for Ftp { }) } + /// FTP replaces the destination silently, so it is checked first; a rename + /// has no atomic precondition to offer, and this is the closest honest one. fn rename( target: Locator, destination: String, ) -> Result { let path = route(&target)?; with_control(|control| { + match probe(control, &destination) { + Ok(_) => return Err(ErrorCode::Conflict), + Err(ErrorCode::NotFound) => {} + Err(error) => return Err(error), + } control.expect(&format!("RNFR {path}"))?; - control.expect(&format!("RNTO {destination}")) - })?; - Self::stat(Locator::Path(destination)) + control.expect(&format!("RNTO {destination}"))?; + probe(control, &destination) + }) } /// A file and an empty directory take different commands, and the server diff --git a/src/listing.rs b/src/listing.rs index cb0ba4f..607dc82 100644 --- a/src/listing.rs +++ b/src/listing.rs @@ -1,48 +1,110 @@ -//! Reading an MLSD listing, which is the one machine-readable form FTP has. +//! Reading MLSD listings and MLST replies, the machine-readable forms FTP has. -use memona_plugin_sdk::{Entry, EntryKind, FileInfo}; +use memona_plugin_sdk::{Entry, EntryKind, ErrorCode, FileInfo, ListPage}; -/// One line of an MLSD listing: `fact=value;...; name`. -pub fn machine_entry(line: &str) -> Option { +/// The facts of one `fact=value;...; name` line. +struct Facts<'a> { + info: FileInfo, + name: &'a str, + /// The listed directory itself or its parent, which is never a child. + own: bool, +} + +fn facts(line: &str) -> Option> { let (facts, name) = line.split_once(' ')?; let name = name.trim(); - if name.is_empty() || name == "." || name == ".." { + if name.is_empty() { return None; } let mut kind = EntryKind::Other; let mut size = 0u64; + let mut own = false; for fact in facts.split(';').filter(|fact| !fact.is_empty()) { let (key, value) = fact.split_once('=')?; - match key.to_ascii_lowercase().as_str() { - "type" => { - kind = match value.to_ascii_lowercase().as_str() { - "file" => EntryKind::File, - "dir" | "cdir" | "pdir" => EntryKind::Directory, - _ => EntryKind::Other, - }; - } - "size" => size = value.parse().unwrap_or(0), - _ => {} + if key.eq_ignore_ascii_case("type") { + kind = if value.eq_ignore_ascii_case("file") { + EntryKind::File + } else if value.eq_ignore_ascii_case("dir") { + EntryKind::Directory + } else if value.eq_ignore_ascii_case("cdir") + || value.eq_ignore_ascii_case("pdir") + { + own = true; + EntryKind::Directory + } else { + EntryKind::Other + }; + } else if key.eq_ignore_ascii_case("size") { + size = value.parse().unwrap_or(0); } } - if matches!(kind, EntryKind::Directory) && name.starts_with('.') { - // A listing may repeat the directory itself; its own name is never a child. - return None; - } - Some(Entry { - name: name.to_owned(), + Some(Facts { info: FileInfo { kind, size, revision: None, }, + name, + own, + }) +} + +/// One child of an MLSD listing. +pub fn machine_entry(line: &str) -> Option { + let facts = facts(line)?; + if facts.own || facts.name == "." || facts.name == ".." { + return None; + } + Some(Entry { + name: facts.name.to_owned(), + info: facts.info, + }) +} + +/// The kind and size in an MLST reply, whose one fact line is indented. +pub fn machine_info(reply: &str) -> Option { + reply + .lines() + .find_map(|line| line.strip_prefix(' ')) + .and_then(facts) + .map(|facts| facts.info) +} + +/// One page of a raw MLSD listing, parsed from a byte offset so a large +/// folder costs only the entries asked for. The cursor is the offset of the +/// first unread line, and it is present whenever any line remains. +pub fn page( + listing: &[u8], + start: usize, + limit: usize, +) -> Result { + let at_line_start = + start == 0 || listing.get(start - 1).is_some_and(|byte| *byte == b'\n'); + if start > listing.len() || !at_line_start { + return Err(ErrorCode::InvalidInput); + } + let mut entries = Vec::new(); + let mut offset = start; + while offset < listing.len() && entries.len() < limit { + let rest = &listing[offset..]; + let end = rest + .iter() + .position(|byte| *byte == b'\n') + .map_or(rest.len(), |index| index + 1); + let line = String::from_utf8_lossy(&rest[..end]); + entries.extend(machine_entry(line.trim_end_matches(['\r', '\n']))); + offset += end; + } + Ok(ListPage { + entries, + cursor: (offset < listing.len()).then(|| offset.to_string()), }) } #[cfg(test)] mod tests { - use super::machine_entry; - use memona_plugin_sdk::EntryKind; + use super::{machine_entry, machine_info, page}; + use memona_plugin_sdk::{EntryKind, ErrorCode}; #[test] fn reads_kind_and_size_and_skips_the_directory_itself() @@ -59,9 +121,12 @@ mod tests { assert!(matches!(directory.info.kind, EntryKind::Directory)); assert_eq!(directory.info.size, 0); - // A name of its own or its parent is not a child of the listing. + // A name of its own or its parent is not a child of the listing, however + // the server spells it, while a hidden directory is an ordinary child. assert!(machine_entry("type=cdir; .").is_none()); assert!(machine_entry("type=pdir; ..").is_none()); + assert!(machine_entry("type=cdir; /papers").is_none()); + assert!(machine_entry("type=dir; .git").is_some()); assert!(machine_entry("type=file;size=1; ").is_none()); assert!(machine_entry("no-space-anywhere").is_none()); Ok(()) @@ -75,4 +140,66 @@ mod tests { assert!(matches!(entry.info.kind, EntryKind::Other)); Ok(()) } + + #[test] + fn an_mlst_reply_names_the_kind_without_a_listing() -> Result<(), &'static str> + { + let file = machine_info( + "250-Begin\n type=file;size=15;modify=20260925063332; /small.txt\n250 End.", + ) + .ok_or("a file")?; + assert!(matches!(file.kind, EntryKind::File)); + assert_eq!(file.size, 15); + // The root is its own parent to some servers, and still a directory. + let root = machine_info("250-Begin\n type=pdir;sizd=4096; /\n250 End.") + .ok_or("root")?; + assert!(matches!(root.kind, EntryKind::Directory)); + assert!(machine_info("250 End.").is_none()); + Ok(()) + } + + #[test] + fn pages_parse_only_from_the_cursor_and_always_point_onward() + -> Result<(), ErrorCode> { + let listing = b"type=cdir; .\r\ntype=file;size=1; a\r\ntype=file;size=2; b\r\ntype=dir; c\r\ntype=pdir; ..\r\n"; + let first = page(listing, 0, 2)?; + let names: Vec<&str> = first + .entries + .iter() + .map(|entry| entry.name.as_str()) + .collect(); + assert_eq!(names, ["a", "b"]); + // The cursor is the byte offset of the next line, so the second page + // never rescans the first one. + let offset: usize = first + .cursor + .ok_or(ErrorCode::Internal)? + .parse() + .map_err(|_| ErrorCode::Internal)?; + assert_eq!( + &listing[..offset], + b"type=cdir; .\r\ntype=file;size=1; a\r\ntype=file;size=2; b\r\n" + ); + let second = page(listing, offset, 1)?; + assert_eq!(second.entries.len(), 1); + assert_eq!(second.entries[0].name, "c"); + // Only the parent line remains: a cursor is still returned, and the page + // it names is empty and final. + let tail: usize = second + .cursor + .ok_or(ErrorCode::Internal)? + .parse() + .map_err(|_| ErrorCode::Internal)?; + let last = page(listing, tail, 1)?; + assert!(last.entries.is_empty() && last.cursor.is_none()); + // A cursor from another listing, or a forged one, is refused. + assert!(matches!(page(listing, 3, 2), Err(ErrorCode::InvalidInput))); + assert!(matches!( + page(listing, listing.len() + 1, 2), + Err(ErrorCode::InvalidInput) + )); + let empty = page(b"", 0, 2)?; + assert!(empty.entries.is_empty() && empty.cursor.is_none()); + Ok(()) + } } diff --git a/src/transfer.rs b/src/transfer.rs new file mode 100644 index 0000000..89f73b8 --- /dev/null +++ b/src/transfer.rs @@ -0,0 +1,117 @@ +//! Filling a pull from a stream that hands back chunks of its own size. + +use memona_plugin_sdk::ErrorCode; + +/// Bytes owed to the caller between pulls. Only a zero-length read means the +/// stream has ended; a short read means nothing. +#[derive(Default)] +pub struct Buffered { + pending: Vec, + ended: bool, +} + +impl Buffered { + /// Whether the stream itself has ended; pending bytes may still remain. + pub fn ended(&self) -> bool { + self.ended + } + + /// Collects up to `limit` bytes, reading until the limit is met or the + /// stream ends. A read that returns more than asked for is kept whole and + /// its surplus waits for the next pull, so no byte is ever dropped. + pub fn pull( + &mut self, + limit: usize, + mut read: impl FnMut(usize) -> Result, ErrorCode>, + ) -> Result, ErrorCode> { + let mut out = if self.pending.len() > limit { + self.pending.drain(..limit).collect() + } else { + std::mem::take(&mut self.pending) + }; + while out.len() < limit && !self.ended { + let room = limit - out.len(); + let chunk = read(room)?; + if chunk.is_empty() { + self.ended = true; + } else if chunk.len() > room { + out.extend_from_slice(&chunk[..room]); + self.pending.extend_from_slice(&chunk[room..]); + } else { + out.extend_from_slice(&chunk); + } + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::Buffered; + use memona_plugin_sdk::ErrorCode; + + /// A stream that answers each read from a script, ignoring the limit. + fn scripted( + chunks: &[&[u8]], + ) -> impl FnMut(usize) -> Result, ErrorCode> { + let mut chunks: Vec> = chunks.iter().map(|c| c.to_vec()).collect(); + chunks.reverse(); + move |_| Ok(chunks.pop().unwrap_or_default()) + } + + #[test] + fn keeps_reading_until_the_limit_is_met() -> Result<(), ErrorCode> { + let mut stream = scripted(&[b"ab", b"cd", b"ef", b""]); + let mut buffered = Buffered::default(); + // Three short reads fill one pull; a short read alone is not the end. + assert_eq!(buffered.pull(6, &mut stream)?, b"abcdef"); + assert!(!buffered.ended()); + assert_eq!(buffered.pull(6, &mut stream)?, b""); + assert!(buffered.ended()); + Ok(()) + } + + #[test] + fn overshoot_waits_for_the_next_pull() -> Result<(), ErrorCode> { + let mut stream = scripted(&[b"abcdefgh", b"ij", b""]); + let mut buffered = Buffered::default(); + assert_eq!(buffered.pull(3, &mut stream)?, b"abc"); + assert_eq!(buffered.pull(3, &mut stream)?, b"def"); + // The remainder of the first chunk comes before anything new is read. + assert_eq!(buffered.pull(3, &mut stream)?, b"ghi"); + assert!(!buffered.ended()); + assert_eq!(buffered.pull(3, &mut stream)?, b"j"); + assert!(buffered.ended()); + assert_eq!(buffered.pull(3, &mut stream)?, b""); + Ok(()) + } + + #[test] + fn the_end_is_only_a_zero_length_read() -> Result<(), ErrorCode> { + let mut stream = scripted(&[b"abcd", b"", b"never"]); + let mut buffered = Buffered::default(); + assert_eq!(buffered.pull(2, &mut stream)?, b"ab"); + assert_eq!(buffered.pull(8, &mut stream)?, b"cd"); + assert!(buffered.ended()); + // Once ended, the stream is never read again. + assert_eq!(buffered.pull(8, &mut stream)?, b""); + Ok(()) + } + + #[test] + fn a_failed_read_is_reported_and_pending_bytes_survive() { + let mut calls = 0; + let mut stream = |_: usize| { + calls += 1; + if calls == 1 { + Ok(b"abcdef".to_vec()) + } else { + Err(ErrorCode::Unavailable) + } + }; + let mut buffered = Buffered::default(); + assert_eq!(buffered.pull(2, &mut stream), Ok(b"ab".to_vec())); + assert_eq!(buffered.pull(8, &mut stream), Err(ErrorCode::Unavailable)); + assert!(!buffered.ended()); + } +}