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 {
                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(PathBuf::from(raw));
    }
    let Some(index) = raw.find(':') else {
        return CopyEndpoint::Local(PathBuf::from(raw));
    };
    let host = &raw[..index];
    let path = &raw[index + 1..];
    if host.is_empty() || host.contains('/') || host.contains('\\') {
        return CopyEndpoint::Local(PathBuf::from(raw));
    }
    CopyEndpoint::Remote {
        host: host.to_string(),
        path: 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() >= 3
        && bytes[0].is_ascii_alphabetic()
        && bytes[1] == b':'
        && (bytes[2] == b'\\' || bytes[2] == b'/')
}

#[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 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"))
        );
    }
}