vtcode_skills/

bundle.rs

//! Skill bundle import/export
//!
//! Supports zip-based packaging of skills for distribution, version management,
//! and inline bundle injection. Follows OpenAI Skills API packaging patterns:
//! - Safe zip extraction with path traversal protection
//! - Size limits (50MB compressed, 25MB per file, 200MB total uncompressed, 500 files max)
//! - Manifest validation after extraction
//! - Versioned storage layout

use anyhow::{Context, Result, bail};
use hashbrown::HashMap;
use std::fs;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};
use vtcode_commons::canonicalize;
use vtcode_commons::safe_archive_relative_path;

/// Maximum compressed bundle size (50 MB)
const MAX_BUNDLE_SIZE: usize = 50 * 1024 * 1024;
/// Maximum uncompressed file size (25 MB)
const MAX_FILE_SIZE: usize = 25 * 1024 * 1024;
/// Maximum total uncompressed size (200 MB)
///
/// `MAX_BUNDLE_SIZE` bounds the *compressed* input, but DEFLATE can expand by
/// ~1000x, so an uncompressed budget must be enforced while writing. Without it
/// a single small bundle can fill the disk (zip bomb).
const MAX_TOTAL_UNCOMPRESSED_SIZE: u64 = 200 * 1024 * 1024;
/// Maximum file count per bundle
const MAX_FILE_COUNT: usize = 500;

/// Whether `total` uncompressed bytes fit within the bundle budget.
///
/// Shared by extraction (incremental, per entry) and post-extraction
/// validation so the two enforcement points cannot drift.
fn within_total_budget(total: u64) -> bool {
    total <= MAX_TOTAL_UNCOMPRESSED_SIZE
}

/// Result of importing a skill bundle
#[derive(Debug, Clone)]
pub struct ImportedSkillInfo {
    name: String,
    version: Option<String>,
    description: String,
    path: PathBuf,
    file_count: usize,
    total_size: u64,
}

/// Export a skill directory to a zip bundle (bytes)
///
/// Walks the skill directory and creates a zip archive.
/// The archive preserves the directory structure with the skill name as root.
pub fn export_skill_bundle(skill_root: &Path) -> Result<Vec<u8>> {
    let skill_md = skill_root.join("SKILL.md");
    if !skill_md.exists() {
        bail!("No SKILL.md found at {}", skill_root.display());
    }

    let mut buf = Vec::new();
    {
        let cursor = std::io::Cursor::new(&mut buf);
        let mut zip_writer = zip::ZipWriter::new(cursor);
        let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated);

        add_dir_to_zip(&mut zip_writer, skill_root, skill_root, options)?;
        zip_writer.finish().context("Failed to finalize zip archive")?;
    }

    info!("Exported skill bundle from {}: {} bytes", skill_root.display(), buf.len());
    Ok(buf)
}

fn add_dir_to_zip<W: Write + std::io::Seek>(
    zip_writer: &mut zip::ZipWriter<W>,
    root: &Path,
    dir: &Path,
    options: zip::write::SimpleFileOptions,
) -> Result<()> {
    for entry in fs::read_dir(dir).with_context(|| format!("reading {}", dir.display()))? {
        let entry = entry?;
        let path = entry.path();
        let rel = path
            .strip_prefix(root)
            .with_context(|| format!("stripping prefix from {}", path.display()))?;

        if path.is_dir() {
            let dir_name = format!("{}/", rel.to_string_lossy());
            zip_writer
                .add_directory(&dir_name, options)
                .with_context(|| format!("adding directory {dir_name}"))?;
            add_dir_to_zip(zip_writer, root, &path, options)?;
        } else {
            let name = rel.to_string_lossy().to_string();
            zip_writer
                .start_file(&name, options)
                .with_context(|| format!("starting file {name}"))?;
            let data = fs::read(&path).with_context(|| format!("reading file {}", path.display()))?;
            zip_writer.write_all(&data)?;
        }
    }
    Ok(())
}

/// Import a skill bundle (zip bytes) into the skill store
///
/// Extracts the zip, validates SKILL.md, and moves to versioned storage.
/// Storage layout: `<dest_store>/<skill-name>/<version>/...`
pub fn import_skill_bundle(zip_bytes: &[u8], dest_store: &Path) -> Result<ImportedSkillInfo> {
    if zip_bytes.len() > MAX_BUNDLE_SIZE {
        bail!("Bundle size {} bytes exceeds maximum {} bytes", zip_bytes.len(), MAX_BUNDLE_SIZE);
    }

    let temp_dir = tempfile::tempdir().context("Failed to create temp directory")?;
    let temp_path = temp_dir.path();

    extract_zip_safely(zip_bytes, temp_path)?;

    let skill_md_path = find_skill_md(temp_path)?;
    let skill_root = skill_md_path.parent().unwrap_or(temp_path);

    validate_extracted_bundle(skill_root)?;

    let (manifest, _instructions) = crate::manifest::parse_skill_content(
        &fs::read_to_string(&skill_md_path).context("Failed to read extracted SKILL.md")?,
    )?;

    let version = manifest.version.clone().unwrap_or_else(|| "0.0.0".to_string());

    let dest_dir = dest_store.join(&manifest.name).join(&version);
    if dest_dir.exists() {
        warn!("Overwriting existing skill version at {}", dest_dir.display());
        fs::remove_dir_all(&dest_dir).context("Failed to remove existing version")?;
    }
    fs::create_dir_all(&dest_dir).context("Failed to create destination directory")?;

    let (file_count, total_size) = copy_dir_recursive(skill_root, &dest_dir)?;

    info!(
        "Imported skill '{}' v{} ({} files, {} bytes) to {}",
        manifest.name,
        version,
        file_count,
        total_size,
        dest_dir.display()
    );

    update_skill_index(dest_store, &manifest.name, &version)?;

    Ok(ImportedSkillInfo {
        name: manifest.name,
        version: Some(version),
        description: manifest.description,
        path: dest_dir,
        file_count,
        total_size,
    })
}

/// Import a skill from base64-encoded zip (inline bundle)
pub fn import_inline_bundle(base64_data: &str, dest_store: &Path) -> Result<ImportedSkillInfo> {
    use base64::Engine;
    let bytes = base64::engine::general_purpose::STANDARD
        .decode(base64_data)
        .context("Failed to decode base64 bundle")?;
    import_skill_bundle(&bytes, dest_store)
}

/// Safely extract a zip archive with path-traversal and size protections
fn extract_zip_safely(zip_bytes: &[u8], dest: &Path) -> Result<()> {
    let cursor = std::io::Cursor::new(zip_bytes);
    let mut archive = zip::ZipArchive::new(cursor).context("Failed to open zip archive")?;

    if archive.len() > MAX_FILE_COUNT {
        bail!("Zip contains {} entries, exceeds maximum {}", archive.len(), MAX_FILE_COUNT);
    }

    let mut total_uncompressed: u64 = 0;

    for i in 0..archive.len() {
        let mut file = archive.by_index(i).with_context(|| format!("reading zip entry {i}"))?;
        let raw_name = file
            .name()
            .with_context(|| format!("zip entry {i} has a non-UTF-8 name"))?
            .into_owned();

        let safe_name = safe_archive_relative_path(Path::new(&raw_name))
            .with_context(|| format!("Path traversal detected in zip entry: {raw_name}"))?;
        if safe_name.as_os_str().is_empty() {
            continue;
        }

        let out_path = dest.join(safe_name);

        if file.is_dir() {
            fs::create_dir_all(&out_path).with_context(|| format!("creating dir {}", out_path.display()))?;
        } else {
            if file.size() > MAX_FILE_SIZE as u64 {
                bail!("Zip entry '{}' ({} bytes) exceeds maximum {} bytes", raw_name, file.size(), MAX_FILE_SIZE);
            }

            if let Some(parent) = out_path.parent() {
                fs::create_dir_all(parent)?;
            }

            let mut out_file =
                fs::File::create(&out_path).with_context(|| format!("creating file {}", out_path.display()))?;
            // Bound by bytes actually written: the archive's declared `file.size()`
            // cannot be trusted, and the cumulative budget stops zip bombs.
            let written = copy_bounded(&mut file, &mut out_file, MAX_FILE_SIZE as u64)
                .with_context(|| format!("writing file {}", out_path.display()))?;
            total_uncompressed += written;
            if !within_total_budget(total_uncompressed) {
                bail!(
                    "Zip uncompressed size {total_uncompressed} bytes exceeds maximum {MAX_TOTAL_UNCOMPRESSED_SIZE} bytes"
                );
            }
        }
    }

    Ok(())
}

/// Copy at most `limit` bytes from `reader` into `writer`.
///
/// Fails rather than silently truncating when the reader yields more than
/// `limit`, so untrusted input whose declared size lies is still capped by the
/// bytes actually written.
fn copy_bounded<R: Read, W: Write>(reader: &mut R, writer: &mut W, limit: u64) -> Result<u64> {
    let mut bounded = reader.take(limit.saturating_add(1));
    let written = std::io::copy(&mut bounded, writer)?;
    if written > limit {
        bail!("entry exceeds maximum {limit} bytes");
    }
    Ok(written)
}

/// Find SKILL.md in extracted directory (handles nested structures)
fn find_skill_md(dir: &Path) -> Result<PathBuf> {
    let direct = dir.join("SKILL.md");
    if direct.exists() {
        return Ok(direct);
    }
    let direct_lower = dir.join("skill.md");
    if direct_lower.exists() {
        return Ok(direct_lower);
    }

    for entry in fs::read_dir(dir).context("Failed to read extracted directory")? {
        let entry = entry?;
        if entry.path().is_dir() {
            let nested = entry.path().join("SKILL.md");
            if nested.exists() {
                return Ok(nested);
            }
            let nested_lower = entry.path().join("skill.md");
            if nested_lower.exists() {
                return Ok(nested_lower);
            }
        }
    }

    bail!("No SKILL.md found in extracted bundle")
}

/// Validate extracted bundle for security
fn validate_extracted_bundle(skill_root: &Path) -> Result<()> {
    let mut file_count = 0u64;
    let mut total_size = 0u64;

    validate_dir_recursive(skill_root, skill_root, &mut file_count, &mut total_size)?;

    if file_count > MAX_FILE_COUNT as u64 {
        bail!("Bundle contains {file_count} files, exceeds maximum {MAX_FILE_COUNT}");
    }

    if !within_total_budget(total_size) {
        bail!("Bundle uncompressed size {total_size} bytes exceeds maximum {MAX_TOTAL_UNCOMPRESSED_SIZE} bytes");
    }

    Ok(())
}

fn validate_dir_recursive(root: &Path, dir: &Path, file_count: &mut u64, total_size: &mut u64) -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();

        if path.is_symlink() {
            bail!("Symlinks not allowed in skill bundles: {}", path.display());
        }

        let canonical = canonicalize(&path).unwrap_or_else(|_| path.clone());
        let root_canonical = canonicalize(root).unwrap_or_else(|_| root.to_path_buf());
        if !canonical.starts_with(&root_canonical) {
            bail!("Path traversal detected: {} escapes bundle root", path.display());
        }

        if path.is_dir() {
            validate_dir_recursive(root, &path, file_count, total_size)?;
        } else {
            *file_count += 1;
            let size = entry.metadata()?.len();
            if size > MAX_FILE_SIZE as u64 {
                bail!("File {} ({size} bytes) exceeds maximum {MAX_FILE_SIZE} bytes", path.display(),);
            }
            *total_size += size;
        }
    }
    Ok(())
}

/// Copy directory recursively, returning (file_count, total_bytes)
fn copy_dir_recursive(src: &Path, dst: &Path) -> Result<(usize, u64)> {
    let mut count = 0usize;
    let mut size = 0u64;

    for entry in fs::read_dir(src)? {
        let entry = entry?;
        let src_path = entry.path();
        let file_name = entry.file_name();
        let dst_path = dst.join(&file_name);

        if src_path.is_dir() {
            fs::create_dir_all(&dst_path)?;
            let (c, s) = copy_dir_recursive(&src_path, &dst_path)?;
            count += c;
            size += s;
        } else {
            fs::copy(&src_path, &dst_path)
                .with_context(|| format!("failed to copy {} to {}", src_path.display(), dst_path.display()))?;
            count += 1;
            size += entry
                .metadata()
                .with_context(|| format!("failed to stat {}", src_path.display()))?
                .len();
        }
    }

    Ok((count, size))
}

/// Skill store index for tracking versions
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
pub struct SkillStoreIndex {
    skills: HashMap<String, SkillVersionIndex>,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SkillVersionIndex {
    latest_version: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    default_version: Option<String>,
    versions: Vec<String>,
}

/// Update the skill index after import
fn update_skill_index(store_path: &Path, skill_name: &str, version: &str) -> Result<()> {
    let index_path = store_path.join("index.json");

    let mut index: SkillStoreIndex = if index_path.exists() {
        let content = fs::read_to_string(&index_path)?;
        serde_json::from_str(&content).unwrap_or_default()
    } else {
        SkillStoreIndex::default()
    };

    let entry = index.skills.entry(skill_name.to_string()).or_insert_with(|| SkillVersionIndex {
        latest_version: version.to_string(),
        default_version: None,
        versions: Vec::new(),
    });

    if !entry.versions.contains(&version.to_string()) {
        entry.versions.push(version.to_string());
    }
    entry.latest_version = version.to_string();

    fs::create_dir_all(store_path)
        .with_context(|| format!("failed to create store dir at {}", store_path.display()))?;
    let index_json =
        serde_json::to_vec_pretty(&index).with_context(|| format!("failed to serialize index for {skill_name}"))?;
    fs::write(&index_path, index_json).with_context(|| format!("failed to write index at {}", index_path.display()))?;

    debug!("Updated skill index at {}", index_path.display());
    Ok(())
}

/// Load the skill store index
pub fn load_skill_index(store_path: &Path) -> Result<SkillStoreIndex> {
    let index_path = store_path.join("index.json");
    if !index_path.exists() {
        return Ok(SkillStoreIndex::default());
    }
    let content = fs::read_to_string(&index_path)?;
    serde_json::from_str(&content).context("Failed to parse skill store index")
}

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

    #[test]
    fn test_skill_store_index_default() {
        let index = SkillStoreIndex::default();
        assert!(index.skills.is_empty());
    }

    #[test]
    fn test_skill_store_index_roundtrip() {
        let mut index = SkillStoreIndex::default();
        index.skills.insert(
            "test-skill".to_string(),
            SkillVersionIndex {
                latest_version: "1.0.0".to_string(),
                default_version: Some("1.0.0".to_string()),
                versions: vec!["0.9.0".to_string(), "1.0.0".to_string()],
            },
        );
        let json = serde_json::to_string(&index).expect("serialize");
        let parsed: SkillStoreIndex = serde_json::from_str(&json).expect("deserialize");
        assert_eq!(parsed.skills["test-skill"].latest_version, "1.0.0");
        assert_eq!(parsed.skills["test-skill"].versions.len(), 2);
    }

    #[test]
    fn test_bundle_size_limit() {
        let oversized = vec![0u8; MAX_BUNDLE_SIZE + 1];
        let temp = tempfile::tempdir().expect("tempdir");
        let result = import_skill_bundle(&oversized, temp.path());
        assert!(result.is_err());
        assert!(result.expect_err("should fail").to_string().contains("exceeds maximum"));
    }

    /// Build an uncompressed zip archive from `(name, contents)` entries.
    fn zip_bytes(entries: &[(&str, &[u8])]) -> Vec<u8> {
        let mut buf = Vec::new();
        {
            let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf));
            let options = zip::write::SimpleFileOptions::default();
            for (name, data) in entries {
                writer.start_file(*name, options).expect("start_file");
                writer.write_all(data).expect("write_all");
            }
            writer.finish().expect("finish");
        }
        buf
    }

    #[test]
    fn copy_bounded_accepts_exactly_the_limit() {
        let data = vec![7u8; 10];
        let mut out = Vec::new();
        let written = copy_bounded(&mut data.as_slice(), &mut out, 10).expect("exact limit is allowed");
        assert_eq!(written, 10);
        assert_eq!(out.len(), 10);
    }

    #[test]
    fn copy_bounded_rejects_one_byte_over_the_limit() {
        let data = vec![7u8; 11];
        let mut out = Vec::new();
        let err = copy_bounded(&mut data.as_slice(), &mut out, 10).expect_err("over limit must fail");
        assert!(err.to_string().contains("exceeds maximum"), "unexpected error: {err}");
    }

    #[test]
    fn within_total_budget_boundaries() {
        assert!(within_total_budget(0));
        assert!(within_total_budget(MAX_TOTAL_UNCOMPRESSED_SIZE));
        assert!(!within_total_budget(MAX_TOTAL_UNCOMPRESSED_SIZE + 1));
    }

    #[test]
    fn extract_rejects_parent_directory_traversal() {
        let bytes = zip_bytes(&[("../evil.txt", b"pwned")]);
        let temp = tempfile::tempdir().expect("tempdir");
        let err = extract_zip_safely(&bytes, temp.path()).expect_err("traversal must fail");
        assert!(err.to_string().contains("Path traversal"), "unexpected error: {err}");
        assert!(!temp.path().join("../evil.txt").exists());
    }

    #[test]
    fn extract_rejects_absolute_path_entries() {
        let bytes = zip_bytes(&[("/abs.txt", b"pwned")]);
        let temp = tempfile::tempdir().expect("tempdir");
        let err = extract_zip_safely(&bytes, temp.path()).expect_err("absolute path must fail");
        assert!(err.to_string().contains("Path traversal"), "unexpected error: {err}");
    }

    #[test]
    fn extract_rejects_mid_path_traversal() {
        let bytes = zip_bytes(&[("skills/../../evil.txt", b"pwned")]);
        let temp = tempfile::tempdir().expect("tempdir");
        let err = extract_zip_safely(&bytes, temp.path()).expect_err("mid-path traversal must fail");
        assert!(err.to_string().contains("Path traversal"), "unexpected error: {err}");
        assert!(!temp.path().join("evil.txt").exists());
    }

    #[test]
    fn extract_accepts_dot_dot_inside_a_single_component() {
        // `..` inside one component is a legitimate file name, not traversal;
        // the substring check this replaced used to reject it.
        let bytes = zip_bytes(&[("v1..2/README.md", b"body")]);
        let temp = tempfile::tempdir().expect("tempdir");
        extract_zip_safely(&bytes, temp.path()).expect("dot-dot component must extract");

        assert_eq!(fs::read(temp.path().join("v1..2/README.md")).expect("read"), b"body");
    }

    #[test]
    fn import_round_trips_an_exported_bundle() {
        let source = tempfile::tempdir().expect("source tempdir");
        let skill_dir = source.path().join("roundtrip-skill");
        fs::create_dir_all(&skill_dir).expect("create skill dir");
        fs::write(
            skill_dir.join("SKILL.md"),
            "---\nname: roundtrip-skill\ndescription: Round trip test skill\nversion: 1.2.3\n---\n# Body\n",
        )
        .expect("write SKILL.md");

        let bytes = export_skill_bundle(&skill_dir).expect("export");

        let store = tempfile::tempdir().expect("store tempdir");
        let imported = import_skill_bundle(&bytes, store.path()).expect("import");
        assert_eq!(imported.name, "roundtrip-skill");
        assert!(imported.file_count >= 1, "expected at least SKILL.md to be copied");
        assert!(imported.path.join("SKILL.md").exists(), "SKILL.md must land in the versioned store");
    }
}