use std::collections::HashMap;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use crate::fs::Entry;
use crate::fs::sftp::SftpFs;
use crate::sftp::wire::{
Attrs, CLOSE, DATA, Dec, Enc, FXF_CREAT, HANDLE, INIT, MKDIR, NAME, OPEN, OPENDIR, READ,
READDIR, REALPATH, STATUS, VERSION, WRITE,
};
use crate::sftp::{read_frame, write_frame};
const SSH_FX_OK: u32 = 0;
const SSH_FX_EOF: u32 = 1;
const SSH_FX_NO_SUCH_FILE: u32 = 2;
pub fn file_attrs(size: u64, mtime: u32) -> Attrs {
Attrs {
size: Some(size),
permissions: Some(0o100644),
mtime: Some(mtime),
..Attrs::default()
}
}
pub fn dir_attrs() -> Attrs {
Attrs {
size: Some(4096),
permissions: Some(0o040755),
mtime: Some(1),
..Attrs::default()
}
}
pub fn symlink_attrs() -> Attrs {
Attrs {
size: Some(11),
permissions: Some(0o120777),
mtime: Some(1),
..Attrs::default()
}
}
#[derive(Default)]
pub struct FakeRemote {
dirs: HashMap<String, Vec<Entry>>,
files: HashMap<String, Vec<u8>>,
reached_as: Option<String>,
refuses: HashMap<String, u32>,
home: Option<String>,
}
impl FakeRemote {
pub fn new() -> Self {
Self::default()
}
pub fn home(mut self, path: &str) -> Self {
self.home = Some(path.to_string());
self
}
pub fn refuses_listing(mut self, path: &str, status: u32) -> Self {
self.refuses.insert(path.to_string(), status);
self
}
pub fn dir(mut self, path: &str, entries: Vec<(&str, Attrs)>) -> Self {
self.dirs.insert(
path.to_string(),
entries
.into_iter()
.map(|(name, attrs)| Entry {
name: name.to_string(),
attrs,
owner: None,
})
.collect(),
);
self
}
pub fn reached_as(mut self, who: &str) -> Self {
self.reached_as = Some(who.to_string());
self
}
pub fn owner(mut self, path: &str, who: &str) -> Self {
let (parent, name) = path.rsplit_once('/').expect("an owner needs a full path");
let entry = self
.dirs
.get_mut(parent)
.and_then(|entries| entries.iter_mut().find(|e| e.name == name))
.unwrap_or_else(|| panic!("no listing entry for {path}; declare it with dir() first"));
entry.owner = Some(who.to_string());
self
}
pub fn file(mut self, path: &str, body: &[u8]) -> Self {
self.files.insert(path.to_string(), body.to_vec());
self
}
pub async fn spawn(self) -> SftpFs {
let (client, server) = tokio::io::duplex(1 << 20);
let (cr, cw) = tokio::io::split(client);
let (sr, sw) = tokio::io::split(server);
tokio::spawn(serve(self, sr, sw));
SftpFs::over(cw, cr)
.await
.expect("handshake with the in-memory remote")
}
}
fn handle_for(kind: char, serial: u32, path: &str) -> Vec<u8> {
format!("{kind}{serial}:{path}").into_bytes()
}
fn path_of(handle: &[u8]) -> Option<String> {
let s = String::from_utf8(handle.to_vec()).ok()?;
let (_kind_and_serial, path) = s.split_once(':')?;
Some(path.to_string())
}
async fn serve<R, W>(mut remote: FakeRemote, mut r: R, mut w: W)
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let (kind, _) = read_frame(&mut r).await.expect("init frame");
assert_eq!(kind, INIT, "first frame must be SSH_FXP_INIT");
write_frame(&mut w, VERSION, &Enc::new().u32(3).done())
.await
.expect("version");
w.flush().await.expect("flush version");
let mut serial = 0u32;
let mut drained: Vec<Vec<u8>> = Vec::new();
while let Ok((kind, payload)) = read_frame(&mut r).await {
let mut d = Dec::new(&payload);
let id = d.u32().expect("request id");
let (out_kind, body) = match kind {
OPENDIR => {
let path = utf8(d.str().expect("opendir path"));
if let Some(&code) = remote.refuses.get(&path) {
(STATUS, status(id, code, "refused"))
} else if remote.dirs.contains_key(&path) {
serial += 1;
(
HANDLE,
Enc::new()
.u32(id)
.str(&handle_for('D', serial, &path))
.done(),
)
} else {
(STATUS, status(id, SSH_FX_NO_SUCH_FILE, "no such directory"))
}
}
READDIR => {
let handle = d.str().expect("readdir handle").to_vec();
let path = path_of(&handle).expect("readdir handle shape");
if drained.contains(&handle) {
(STATUS, status(id, SSH_FX_EOF, "eof"))
} else {
drained.push(handle);
let entries = remote.dirs.get(&path).expect("listed dir exists");
(NAME, names(id, entries))
}
}
OPEN => {
let path = utf8(d.str().expect("open path"));
let flags = d.u32().expect("open flags");
let creating = flags & FXF_CREAT != 0;
if creating {
remote.files.entry(path.clone()).or_default();
}
if remote.files.contains_key(&path) {
serial += 1;
(
HANDLE,
Enc::new()
.u32(id)
.str(&handle_for('F', serial, &path))
.done(),
)
} else {
(STATUS, status(id, SSH_FX_NO_SUCH_FILE, "no such file"))
}
}
READ => {
let handle = d.str().expect("read handle").to_vec();
let path = path_of(&handle).expect("read handle shape");
let offset = d.u64().expect("read offset") as usize;
let want = d.u32().expect("read length") as usize;
let body = remote.files.get(&path).expect("opened file exists");
if offset >= body.len() {
(STATUS, status(id, SSH_FX_EOF, "eof"))
} else {
let end = offset.saturating_add(want).min(body.len());
(DATA, Enc::new().u32(id).str(&body[offset..end]).done())
}
}
WRITE => {
let handle = d.str().expect("write handle").to_vec();
let path = path_of(&handle).expect("write handle shape");
d.u64().expect("write offset");
let data = d.str().expect("write data").to_vec();
remote
.files
.entry(path.clone())
.or_default()
.extend_from_slice(&data);
if let Some((parent, name)) = path.rsplit_once('/') {
let size = remote.files.get(&path).map_or(0, Vec::len) as u64;
let entries = remote.dirs.entry(parent.to_string()).or_default();
match entries.iter_mut().find(|e| e.name == name) {
Some(e) => e.attrs.size = Some(size),
None => entries.push(Entry {
name: name.to_string(),
attrs: file_attrs(size, 1),
owner: remote.reached_as.clone(),
}),
}
}
(STATUS, status(id, SSH_FX_OK, "ok"))
}
REALPATH => {
d.str().expect("realpath path");
match &remote.home {
Some(home) => (
NAME,
names(
id,
&[Entry {
name: home.clone(),
attrs: dir_attrs(),
owner: remote.reached_as.clone(),
}],
),
),
None => (
STATUS,
status(id, 4, "this remote was not given a home directory"),
),
}
}
MKDIR => {
let path = utf8(d.str().expect("mkdir path"));
remote.dirs.entry(path).or_default();
(STATUS, status(id, SSH_FX_OK, "ok"))
}
CLOSE => {
d.str().expect("close handle");
(STATUS, status(id, SSH_FX_OK, "ok"))
}
other => panic!("in-memory remote got unexpected request type {other}"),
};
write_frame(&mut w, out_kind, &body).await.expect("reply");
w.flush().await.expect("flush reply");
}
}
fn utf8(bytes: &[u8]) -> String {
String::from_utf8_lossy(bytes).into_owned()
}
fn status(id: u32, code: u32, message: &str) -> Vec<u8> {
Enc::new()
.u32(id)
.u32(code)
.str(message.as_bytes())
.str(b"")
.done()
}
const WRITTEN_ATTRS: u32 = 0x0000_0001 | 0x0000_0004 | 0x0000_0008;
fn longname(e: &Entry) -> String {
match &e.owner {
Some(who) => format!(
"{} 1 {who} {who} {:>8} Jan 1 00:00 {}",
mode_column(&e.attrs),
e.attrs.size.unwrap_or(0),
e.name
),
None => e.name.clone(),
}
}
fn mode_column(attrs: &Attrs) -> &'static str {
if attrs.is_dir() {
"drwxr-xr-x"
} else if attrs.is_symlink() {
"lrwxrwxrwx"
} else {
"-rw-r--r--"
}
}
fn names(id: u32, entries: &[Entry]) -> Vec<u8> {
let mut enc = Enc::new().u32(id).u32(entries.len() as u32);
for e in entries {
enc = enc
.str(e.name.as_bytes())
.str(longname(e).as_bytes())
.u32(WRITTEN_ATTRS)
.u64(e.attrs.size.unwrap_or(0))
.u32(e.attrs.permissions.unwrap_or(0o100644))
.u32(e.attrs.atime.unwrap_or(0))
.u32(e.attrs.mtime.unwrap_or(0));
}
enc.done()
}