megacode-agent-sdk 0.31.0

//! The release archive reader (spec slice-ar R12.2): gzip, then tar; the one
//! regular file `harness` (or `./harness`) is the executable. Pax headers
//! (`x`, `g`) and AppleDouble members (`._*`) are skipped; any other member
//! type, a second regular file, or a path with `..` is an error.

use std::io::Read;

/// The largest executable accepted.
pub(crate) const MAX_EXECUTABLE: u64 = 256 * 1024 * 1024;

/// What the skipped members (pax headers, AppleDouble files, empty
/// folders) may add up to.
const MAX_SKIPPED: u64 = 16 * 1024 * 1024;

const BLOCK: usize = 512;

/// The executable's bytes from a gzipped tar.
pub(crate) fn executable(gz: &[u8]) -> Result<Vec<u8>, String> {
    let mut input = flate2::read::GzDecoder::new(gz);
    let mut found: Option<Vec<u8>> = None;
    let mut skipped = 0u64;
    loop {
        let mut header = [0u8; BLOCK];
        if let Err(e) = input.read_exact(&mut header) {
            return Err(format!("the archive ends early: {e}"));
        }
        if header.iter().all(|&b| b == 0) {
            break;
        }
        let member = Header::parse(&header)?;
        if member.path.split('/').any(|part| part == "..") {
            return Err(format!("the archive holds a path with ..: {}", member.path));
        }
        let name = member.path.strip_prefix("./").unwrap_or(&member.path);
        let base = name.rsplit('/').next().unwrap_or(name);
        let skip = match member.kind {
            b'x' | b'g' => true,
            b'0' | 0 if base.starts_with("._") => true,
            b'5' if member.size == 0 => true,
            b'0' | 0 if name == "harness" => {
                if found.is_some() {
                    return Err("the archive holds a second harness".into());
                }
                if member.size > MAX_EXECUTABLE {
                    return Err(format!(
                        "the archive's harness is {} bytes, over 256 MiB",
                        member.size
                    ));
                }
                let mut body = Vec::with_capacity(member.size as usize);
                read_body(&mut input, member.size, Some(&mut body))?;
                found = Some(body);
                false
            }
            b'0' | 0 => {
                return Err(format!(
                    "the archive holds a regular file other than harness: {}",
                    member.path
                ));
            }
            kind => {
                return Err(format!(
                    "the archive holds a member of type {:?} that is not a regular file: {}",
                    char::from(kind),
                    member.path
                ));
            }
        };
        if skip {
            skipped += member.size;
            if skipped > MAX_SKIPPED {
                return Err("the archive's other members are over 16 MiB".into());
            }
            read_body(&mut input, member.size, None)?;
        }
    }
    found.ok_or_else(|| "the archive holds no harness".into())
}

/// Reads a member's body of `size` bytes and its padding to the block;
/// keeps the body in `keep` when given.
fn read_body(
    input: &mut impl Read,
    size: u64,
    mut keep: Option<&mut Vec<u8>>,
) -> Result<(), String> {
    let padded = size.div_ceil(BLOCK as u64) * BLOCK as u64;
    let mut buf = [0u8; 64 * 1024];
    let mut done = 0u64;
    while done < padded {
        let want = ((padded - done) as usize).min(buf.len());
        input
            .read_exact(&mut buf[..want])
            .map_err(|e| format!("the archive ends early: {e}"))?;
        if let Some(out) = keep.as_deref_mut() {
            let body_left = size.saturating_sub(done) as usize;
            out.extend_from_slice(&buf[..want.min(body_left)]);
        }
        done += want as u64;
    }
    Ok(())
}

/// The fields of a tar header this reader uses.
struct Header {
    path: String,
    kind: u8,
    size: u64,
}

impl Header {
    fn parse(h: &[u8; BLOCK]) -> Result<Header, String> {
        let stored = octal(&h[148..156]).ok_or("the archive has a bad header checksum")?;
        let sum: u64 = h
            .iter()
            .enumerate()
            .map(|(i, &b)| {
                if (148..156).contains(&i) {
                    32
                } else {
                    u64::from(b)
                }
            })
            .sum();
        if sum != stored {
            return Err("the archive has a bad header checksum".into());
        }
        let size = octal(&h[124..136]).ok_or("the archive has a bad member size")?;
        let name = text(&h[0..100]);
        let prefix = if &h[257..262] == b"ustar" {
            text(&h[345..500])
        } else {
            String::new()
        };
        let path = if prefix.is_empty() {
            name
        } else {
            format!("{prefix}/{name}")
        };
        Ok(Header {
            path,
            kind: h[156],
            size,
        })
    }
}

/// A NUL- or space-ended octal field; `None` when it is not one (a
/// base-256 size is over any limit here, so it is refused as malformed).
fn octal(field: &[u8]) -> Option<u64> {
    let digits: Vec<u8> = field
        .iter()
        .copied()
        .skip_while(|&b| b == b' ')
        .take_while(|&b| b != 0 && b != b' ')
        .collect();
    if digits.is_empty() || digits.len() > 21 {
        return if digits.is_empty() { Some(0) } else { None };
    }
    digits.iter().try_fold(0u64, |n, &d| {
        (b'0'..=b'7')
            .contains(&d)
            .then(|| n * 8 + u64::from(d - b'0'))
    })
}

fn text(field: &[u8]) -> String {
    let end = field.iter().position(|&b| b == 0).unwrap_or(field.len());
    String::from_utf8_lossy(&field[..end]).into_owned()
}

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

    fn header(name: &str, kind: u8, size: u64) -> [u8; BLOCK] {
        let mut h = [0u8; BLOCK];
        h[..name.len()].copy_from_slice(name.as_bytes());
        h[124..135].copy_from_slice(format!("{size:011o}").as_bytes());
        h[156] = kind;
        h[257..263].copy_from_slice(b"ustar\0");
        h[148..156].copy_from_slice(b"        ");
        let sum: u32 = h.iter().map(|&b| u32::from(b)).sum();
        h[148..156].copy_from_slice(format!("{sum:06o}\0 ").as_bytes());
        h
    }

    fn gz(blocks: &[&[u8]]) -> Vec<u8> {
        let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
        for b in blocks {
            e.write_all(b).unwrap();
        }
        e.write_all(&[0u8; 1024]).unwrap();
        e.finish().unwrap()
    }

    #[test]
    fn an_executable_over_256_mib_is_refused_from_its_header() {
        let h = header("harness", b'0', MAX_EXECUTABLE + 1);
        let err = executable(&gz(&[&h])).unwrap_err();
        assert!(err.contains("over 256 MiB"), "{err}");
    }

    #[test]
    fn a_corrupt_header_is_refused() {
        let mut h = header("harness", b'0', 0);
        h[0] = b'H';
        let err = executable(&gz(&[&h])).unwrap_err();
        assert!(err.contains("checksum"), "{err}");
    }

    #[test]
    fn a_folder_with_content_or_a_device_is_refused() {
        for (kind, size) in [(b'5', 1), (b'3', 0), (b'4', 0), (b'6', 0)] {
            let h = header("harness", kind, size);
            let err = executable(&gz(&[&h, &[0u8; BLOCK]])).unwrap_err();
            assert!(err.contains("not a regular file"), "{err}");
        }
    }

    #[test]
    fn the_body_is_read_exactly() {
        let mut body = vec![7u8; 700];
        body.resize(1024, 0);
        let h = header("harness", b'0', 700);
        assert_eq!(executable(&gz(&[&h, &body])).unwrap(), vec![7u8; 700]);
    }
}