dosh/src/file_transfer.rs

use anyhow::{Context, Result, anyhow, bail};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

pub const FILE_STREAM_SENTINEL: &str = "@dosh-file";
pub const FRAME_MAX_LEN: usize = 1024 * 1024;
pub const CHUNK_SIZE: usize = 32 * 1024;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum FileRequest {
    Stat {
        path: String,
    },
    List {
        path: String,
    },
    Mkdir {
        path: String,
        mode: Option<u32>,
    },
    Readlink {
        path: String,
    },
    Symlink {
        path: String,
        target: String,
        overwrite: bool,
    },
    Remove {
        path: String,
        recursive: bool,
    },
    PutStart {
        path: String,
        size: u64,
        mode: Option<u32>,
        modified_secs: Option<u64>,
        overwrite: bool,
        resume: bool,
    },
    PutChunk {
        offset: u64,
        bytes: Vec<u8>,
    },
    PutFinish {
        sha256: [u8; 32],
    },
    Get {
        path: String,
        offset: u64,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum FileResponse {
    Ok,
    Resume {
        offset: u64,
        prefix_sha256: Option<[u8; 32]>,
    },
    Stat {
        meta: FileMeta,
    },
    List {
        entries: Vec<FileEntry>,
    },
    LinkTarget {
        target: String,
    },
    Start {
        meta: FileMeta,
        offset: u64,
        prefix_sha256: Option<[u8; 32]>,
    },
    Chunk {
        offset: u64,
        bytes: Vec<u8>,
    },
    Done {
        sha256: [u8; 32],
        bytes: u64,
    },
    Error {
        message: String,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct FileMeta {
    pub path: String,
    pub kind: FileKind,
    pub len: u64,
    pub mode: Option<u32>,
    pub modified_secs: Option<u64>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum FileKind {
    File,
    Directory,
    Symlink,
    Other,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct FileEntry {
    pub name: String,
    pub meta: FileMeta,
}

#[derive(Default)]
pub struct FrameDecoder {
    buf: Vec<u8>,
}

impl FrameDecoder {
    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {
        self.buf.extend_from_slice(bytes);
        let mut frames = Vec::new();
        loop {
            if self.buf.len() < 4 {
                break;
            }
            let len = u32::from_be_bytes(self.buf[..4].try_into().unwrap()) as usize;
            if len > FRAME_MAX_LEN {
                self.buf.clear();
                bail!("file protocol frame too large: {len} bytes");
            }
            if self.buf.len() < 4 + len {
                break;
            }
            frames.push(self.buf[4..4 + len].to_vec());
            self.buf.drain(..4 + len);
        }
        Ok(frames)
    }
}

pub fn encode_request(request: &FileRequest) -> Result<Vec<u8>> {
    encode_frame(&bincode::serialize(request).context("encode file request")?)
}

pub fn encode_response(response: &FileResponse) -> Result<Vec<u8>> {
    encode_frame(&bincode::serialize(response).context("encode file response")?)
}

pub fn decode_request(frame: &[u8]) -> Result<FileRequest> {
    bincode::deserialize(frame).context("decode file request")
}

pub fn decode_response(frame: &[u8]) -> Result<FileResponse> {
    bincode::deserialize(frame).context("decode file response")
}

pub fn encode_frame(payload: &[u8]) -> Result<Vec<u8>> {
    if payload.len() > FRAME_MAX_LEN {
        bail!("file protocol frame too large: {} bytes", payload.len());
    }
    let mut out = Vec::with_capacity(4 + payload.len());
    out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
    out.extend_from_slice(payload);
    Ok(out)
}

pub fn parse_copy_endpoint(raw: &str) -> CopyEndpoint {
    if looks_like_windows_path(raw) {
        return CopyEndpoint::Local(local_copy_path(raw));
    }
    if let Some((host, path)) = parse_bracketed_remote_endpoint(raw) {
        return CopyEndpoint::Remote { host, path };
    }
    if raw.starts_with('[') {
        return CopyEndpoint::Local(local_copy_path(raw));
    }
    let Some(index) = raw.find(':') else {
        return CopyEndpoint::Local(local_copy_path(raw));
    };
    let host = &raw[..index];
    let path = &raw[index + 1..];
    if host.is_empty() || host.contains('/') || host.contains('\\') {
        return CopyEndpoint::Local(local_copy_path(raw));
    }
    CopyEndpoint::Remote {
        host: host.to_string(),
        path: if path.is_empty() {
            ".".to_string()
        } else {
            path.to_string()
        },
    }
}

fn parse_bracketed_remote_endpoint(raw: &str) -> Option<(String, String)> {
    let rest = raw.strip_prefix('[')?;
    let close = rest.find(']')?;
    let host = &rest[..close];
    let suffix = &rest[close + 1..];
    if host.is_empty() || host.contains('/') || host.contains('\\') || !suffix.starts_with(':') {
        return None;
    }
    let path = &suffix[1..];
    Some((
        host.to_string(),
        if path.is_empty() {
            ".".to_string()
        } else {
            path.to_string()
        },
    ))
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CopyEndpoint {
    Local(PathBuf),
    Remote { host: String, path: String },
}

pub fn remote_join(parent: &str, child: &str) -> String {
    if parent.is_empty() || parent == "." {
        return child.to_string();
    }
    format!("{}/{}", parent.trim_end_matches('/'), child)
}

pub fn remote_basename(path: &str) -> String {
    path.trim_end_matches('/')
        .rsplit('/')
        .next()
        .filter(|value| !value.is_empty())
        .unwrap_or(path)
        .to_string()
}

pub fn clean_remote_path(path: &str, home: &Path) -> Result<PathBuf> {
    if path.as_bytes().contains(&0) {
        bail!("remote path contains NUL");
    }
    let expanded = if path == "~" {
        home.to_path_buf()
    } else if let Some(rest) = path.strip_prefix("~/") {
        home.join(rest)
    } else {
        let raw = Path::new(path);
        if raw.is_absolute() {
            raw.to_path_buf()
        } else {
            home.join(raw)
        }
    };
    Ok(expanded)
}

pub fn ensure_relative_child(path: &Path) -> Result<String> {
    let name = path
        .file_name()
        .and_then(|name| name.to_str())
        .ok_or_else(|| anyhow!("path has no final component: {}", path.display()))?;
    if name.is_empty() || name == "." || name == ".." {
        bail!("invalid path final component: {}", path.display());
    }
    Ok(name.to_string())
}

fn looks_like_windows_path(raw: &str) -> bool {
    let bytes = raw.as_bytes();
    bytes.len() >= 2 && bytes[0].is_ascii_alphabetic() && bytes[1] == b':'
}

fn local_copy_path(raw: &str) -> PathBuf {
    let rest = if raw == "~" {
        Some("")
    } else {
        raw.strip_prefix("~/").or_else(|| raw.strip_prefix("~\\"))
    };
    if let Some(rest) = rest
        && let Some(home) = dirs::home_dir()
    {
        if rest.is_empty() {
            return home;
        }
        return home.join(rest);
    }
    PathBuf::from(raw)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn frames_round_trip_incrementally() {
        let first = encode_response(&FileResponse::Ok).unwrap();
        let second = encode_response(&FileResponse::Error {
            message: "nope".to_string(),
        })
        .unwrap();
        let mut decoder = FrameDecoder::default();
        assert!(decoder.push(&first[..2]).unwrap().is_empty());
        let frames = decoder.push(&first[2..]).unwrap();
        assert_eq!(frames.len(), 1);
        assert_eq!(decode_response(&frames[0]).unwrap(), FileResponse::Ok);
        let frames = decoder.push(&second).unwrap();
        assert_eq!(frames.len(), 1);
    }

    #[test]
    fn oversized_frame_clears_decoder_state() {
        let mut decoder = FrameDecoder::default();
        let too_large = ((FRAME_MAX_LEN + 1) as u32).to_be_bytes();
        assert!(decoder.push(&too_large).is_err());

        let valid = encode_response(&FileResponse::Ok).unwrap();
        let frames = decoder.push(&valid).unwrap();
        assert_eq!(frames.len(), 1);
        assert_eq!(decode_response(&frames[0]).unwrap(), FileResponse::Ok);
    }

    #[test]
    fn copy_endpoint_parses_scp_style_but_not_windows_drive() {
        assert_eq!(
            parse_copy_endpoint("palav:tmp/file"),
            CopyEndpoint::Remote {
                host: "palav".to_string(),
                path: "tmp/file".to_string()
            }
        );
        assert_eq!(
            parse_copy_endpoint("C:\\Users\\palav\\x"),
            CopyEndpoint::Local(PathBuf::from("C:\\Users\\palav\\x"))
        );
        assert_eq!(
            parse_copy_endpoint("C:/Users/palav/x"),
            CopyEndpoint::Local(PathBuf::from("C:/Users/palav/x"))
        );
        assert_eq!(
            parse_copy_endpoint("C:relative\\x"),
            CopyEndpoint::Local(PathBuf::from("C:relative\\x"))
        );
    }

    #[test]
    fn copy_endpoint_parses_bracketed_ipv6_remote_paths() {
        assert_eq!(
            parse_copy_endpoint("[2001:db8::1]:/var/log/syslog"),
            CopyEndpoint::Remote {
                host: "2001:db8::1".to_string(),
                path: "/var/log/syslog".to_string()
            }
        );
        assert_eq!(
            parse_copy_endpoint("[::1]:"),
            CopyEndpoint::Remote {
                host: "::1".to_string(),
                path: ".".to_string()
            }
        );
        assert_eq!(
            parse_copy_endpoint("[bad/host]:path"),
            CopyEndpoint::Local(PathBuf::from("[bad/host]:path"))
        );
        assert_eq!(
            parse_copy_endpoint("[::1]"),
            CopyEndpoint::Local(PathBuf::from("[::1]"))
        );
    }

    #[test]
    fn copy_endpoint_expands_local_home_paths() {
        let Some(home) = dirs::home_dir() else {
            return;
        };
        assert_eq!(parse_copy_endpoint("~"), CopyEndpoint::Local(home.clone()));
        assert_eq!(
            parse_copy_endpoint("~/Downloads/file.txt"),
            CopyEndpoint::Local(home.join("Downloads/file.txt"))
        );
        assert_eq!(
            parse_copy_endpoint("~\\Downloads\\file.txt"),
            CopyEndpoint::Local(home.join("Downloads\\file.txt"))
        );
        assert_eq!(
            parse_copy_endpoint("host:~/Downloads/file.txt"),
            CopyEndpoint::Remote {
                host: "host".to_string(),
                path: "~/Downloads/file.txt".to_string()
            }
        );
    }
}