magi-code 0.63.4

use std::{fs::OpenOptions, io::Read, path::Path};

#[derive(Debug, PartialEq, Eq)]
pub(crate) struct PromptFile {
    pub(crate) text: String,
    pub(crate) bytes: usize,
}

pub(crate) fn read_prompt_file(
    path: &Path,
    limit: u64,
    reject_symlink: bool,
) -> anyhow::Result<PromptFile> {
    #[cfg(not(any(unix, windows)))]
    if reject_symlink && std::fs::symlink_metadata(path)?.file_type().is_symlink() {
        anyhow::bail!("prompt file must not be a symlink: {}", path.display());
    }
    let mut options = OpenOptions::new();
    options.read(true);
    if reject_symlink {
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            options.custom_flags(o_no_follow());
        }
        #[cfg(windows)]
        {
            use std::os::windows::fs::OpenOptionsExt;
            options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT);
        }
    }
    let file = match options.open(path) {
        Ok(file) => file,
        Err(error)
            if reject_symlink
                && (std::fs::symlink_metadata(path)
                    .map(|metadata| metadata.file_type().is_symlink())
                    .unwrap_or(false)
                    || matches!(error.raw_os_error(), Some(40 | 62))) =>
        {
            anyhow::bail!("prompt file must not be a symlink: {}", path.display());
        }
        Err(error) => return Err(error.into()),
    };
    if !file.metadata()?.is_file() {
        anyhow::bail!("prompt file must be a regular file: {}", path.display());
    }
    let read_limit = limit
        .checked_add(1)
        .ok_or_else(|| anyhow::anyhow!("prompt file limit is too large: {limit}"))?;
    let mut bytes = Vec::new();
    file.take(read_limit).read_to_end(&mut bytes)?;
    if bytes.len() as u64 > limit {
        anyhow::bail!(
            "prompt file exceeds limit (exceeding); limit is {limit} bytes: {}",
            path.display()
        );
    }
    let bytes_len = bytes.len();
    let text = String::from_utf8(bytes)
        .map_err(|_| anyhow::anyhow!("prompt file is not valid UTF-8: {}", path.display()))?;
    Ok(PromptFile {
        text,
        bytes: bytes_len,
    })
}

#[cfg(unix)]
fn o_no_follow() -> i32 {
    #[cfg(target_os = "linux")]
    {
        0x20000
    }
    #[cfg(not(target_os = "linux"))]
    {
        0x100
    }
}

#[cfg(windows)]
const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x00200000;

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    #[test]
    fn accepts_exact_limit_and_rejects_limit_plus_one() {
        let temp = TempDir::new().unwrap();
        let path = temp.path().join("prompt.md");
        fs::write(&path, b"1234").unwrap();
        assert_eq!(read_prompt_file(&path, 4, false).unwrap().bytes, 4);
        fs::write(&path, b"12345").unwrap();
        let error = read_prompt_file(&path, 4, false).unwrap_err().to_string();
        assert!(error.contains("exceeds"));
    }

    #[test]
    fn rejects_invalid_utf8_without_echoing_bytes() {
        let temp = TempDir::new().unwrap();
        let path = temp.path().join("prompt.md");
        fs::write(&path, [b'a', 0xff]).unwrap();
        let error = read_prompt_file(&path, 2, false).unwrap_err().to_string();
        assert!(error.contains("not valid UTF-8"));
        assert!(!error.contains("0xff"));
    }

    #[test]
    fn rejects_non_regular_file() {
        let temp = TempDir::new().unwrap();
        let error = read_prompt_file(temp.path(), 64, false)
            .unwrap_err()
            .to_string();
        assert!(error.contains("regular file"), "{error}");
    }

    #[cfg(unix)]
    #[test]
    fn no_follow_open_rejects_symlink_replacement() {
        use std::os::unix::fs::symlink;
        let temp = TempDir::new().unwrap();
        let target = temp.path().join("target.md");
        let link = temp.path().join("link.md");
        fs::write(&target, "secret").unwrap();
        symlink(&target, &link).unwrap();
        let error = read_prompt_file(&link, 64, true).unwrap_err().to_string();
        assert!(
            error.contains("symlink") || error.contains("Too many levels"),
            "{error}"
        );
        assert_eq!(read_prompt_file(&link, 64, false).unwrap().text, "secret");
    }
}