//! What a rootfs image is for, read from its own files (OxygenFlow spec
//! §11.2): the ELF header of its shell. An image is an ext4 filesystem,
//! as `flow rootfs import` makes it; this reads just enough of one to
//! follow a path, symlinks included, to a file's first bytes. It needs no
//! mount and no privileges, so a runner can check an image it fetched and
//! a hub an image it keeps as chunks.
//!
//! Read only, and defensive: a malformed image answers `None`, never a
//! panic or an unbounded read.

use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

/// The files whose architecture is the image's, in order: the shell's,
/// else the first program found.
const PROGRAMS: [&str; 5] = ["bin/sh", "bin/busybox", "usr/bin/env", "bin/bash", "usr/bin/sh"];

const MAGIC: u16 = 0xEF53;
const ROOT_INODE: u32 = 2;
const INCOMPAT_64BIT: u32 = 0x80;
const EXTENTS_FL: u32 = 0x80000;
const INLINE_DATA_FL: u32 = 0x1000_0000;
const EXTENT_MAGIC: u16 = 0xF30A;
/// More symlinks than this on one path is a loop.
const MAX_LINKS: u32 = 40;
/// The largest directory read, in bytes.
const MAX_DIR: u64 = 64 << 20;

/// The architecture an ELF header names (`x86_64`, `aarch64`, `riscv64`,
/// …), as `std::env::consts::ARCH` spells it; `None` for anything else.
pub fn elf_arch(header: &[u8]) -> Option<&'static str> {
    if header.len() < 20 || &header[..4] != b"\x7fELF" {
        return None;
    }
    let wide = header[4] == 2;
    let machine = if header[5] == 1 { u16::from_le_bytes([header[18], header[19]]) } else { u16::from_be_bytes([header[18], header[19]]) };
    Some(match machine {
        0x3E => "x86_64",
        0xB7 => "aarch64",
        0xF3 if wide => "riscv64",
        0xF3 => "riscv32",
        0x03 => "x86",
        0x28 => "arm",
        0x15 => "powerpc64",
        0x16 => "s390x",
        _ => return None,
    })
}

/// The architecture of the ext4 image read through `read` (fill the
/// buffer from that byte offset): its shell's, else the first program of
/// [`PROGRAMS`] it has. `None` when it is not ext4 or has none of them.
pub fn arch(read: &mut dyn FnMut(u64, &mut [u8]) -> std::io::Result<()>) -> Option<&'static str> {
    let fs = Fs::open(read)?;
    PROGRAMS.iter().find_map(|p| {
        let inode = fs.resolve(read, p)?;
        if inode.mode & 0xF000 != 0x8000 {
            return None;
        }
        let mut head = [0u8; 20];
        fs.read_file(read, &inode, &mut head)?;
        elf_arch(&head)
    })
}

/// The architecture of the ext4 image at `path` ([`arch`]).
pub fn arch_of_file(path: &Path) -> Option<&'static str> {
    let mut f = std::fs::File::open(path).ok()?;
    arch(&mut |at: u64, buf: &mut [u8]| {
        f.seek(SeekFrom::Start(at))?;
        f.read_exact(buf)
    })
}

struct Fs {
    block: u64,
    first_data_block: u64,
    inodes_per_group: u32,
    inode_size: u64,
    desc_size: u64,
    wide: bool,
    groups: u32,
}

struct Inode {
    mode: u16,
    size: u64,
    flags: u32,
    block: [u8; 60],
}

fn u16_at(b: &[u8], at: usize) -> u16 {
    u16::from_le_bytes([b[at], b[at + 1]])
}

fn u32_at(b: &[u8], at: usize) -> u32 {
    u32::from_le_bytes([b[at], b[at + 1], b[at + 2], b[at + 3]])
}

type ReadFn<'a> = dyn FnMut(u64, &mut [u8]) -> std::io::Result<()> + 'a;

impl Fs {
    fn open(read: &mut ReadFn) -> Option<Fs> {
        let mut sb = [0u8; 1024];
        read(1024, &mut sb).ok()?;
        if u16_at(&sb, 0x38) != MAGIC {
            return None;
        }
        let log = u32_at(&sb, 0x18);
        if log > 6 {
            return None;
        }
        let block = 1024u64 << log;
        let inode_size = if u32_at(&sb, 0x4C) == 0 { 128 } else { u16_at(&sb, 0x58) as u64 };
        let wide = u32_at(&sb, 0x60) & INCOMPAT_64BIT != 0;
        let desc_size = if wide { (u16_at(&sb, 0xFE) as u64).max(64) } else { 32 };
        let inodes_per_group = u32_at(&sb, 0x28);
        let inodes = u32_at(&sb, 0);
        if !(128..=1024).contains(&inode_size) || desc_size > 1024 || inodes_per_group == 0 {
            return None;
        }
        Some(Fs {
            block,
            first_data_block: u32_at(&sb, 0x14) as u64,
            inodes_per_group,
            inode_size,
            desc_size,
            wide,
            groups: inodes.div_ceil(inodes_per_group),
        })
    }

    fn inode(&self, read: &mut ReadFn, n: u32) -> Option<Inode> {
        let (group, index) = ((n.checked_sub(1)?) / self.inodes_per_group, (n - 1) % self.inodes_per_group);
        if group >= self.groups {
            return None;
        }
        let mut desc = [0u8; 64];
        let at = (self.first_data_block + 1) * self.block + group as u64 * self.desc_size;
        let len = if self.wide { 64 } else { 32 };
        read(at, &mut desc[..len]).ok()?;
        let hi = if self.wide { u32_at(&desc, 0x28) as u64 } else { 0 };
        let table = (hi << 32) | u32_at(&desc, 8) as u64;
        let mut raw = [0u8; 160];
        read(table * self.block + index as u64 * self.inode_size, &mut raw[..128]).ok()?;
        let mut block = [0u8; 60];
        block.copy_from_slice(&raw[0x28..0x28 + 60]);
        Some(Inode {
            mode: u16_at(&raw, 0),
            size: ((u32_at(&raw, 0x6C) as u64) << 32) | u32_at(&raw, 4) as u64,
            flags: u32_at(&raw, 0x20),
            block,
        })
    }

    /// The physical block holding logical block `n` of `inode`; `None` for
    /// a hole or an extent not yet written.
    fn map(&self, read: &mut ReadFn, inode: &Inode, n: u64) -> Option<u64> {
        if inode.flags & EXTENTS_FL != 0 {
            let mut node = inode.block.to_vec();
            for _ in 0..6 {
                if u16_at(&node, 0) != EXTENT_MAGIC {
                    return None;
                }
                let entries = (u16_at(&node, 2) as usize).min((node.len() - 12) / 12);
                let depth = u16_at(&node, 6);
                let entry = |i: usize| &node[12 + 12 * i..24 + 12 * i];
                if depth == 0 {
                    return (0..entries).map(entry).find_map(|e| {
                        let (first, len) = (u32_at(e, 0) as u64, u16_at(e, 4) as u64);
                        let start = ((u16_at(e, 6) as u64) << 32) | u32_at(e, 8) as u64;
                        (len <= 32768 && n >= first && n < first + len).then_some(start + (n - first))
                    });
                }
                let next = (0..entries).map(entry).take_while(|e| u32_at(e, 0) as u64 <= n).last()?;
                let child = ((u16_at(next, 8) as u64) << 32) | u32_at(next, 4) as u64;
                let mut buf = vec![0u8; self.block as usize];
                read(child * self.block, &mut buf).ok()?;
                node = buf;
            }
            return None;
        }
        // The old block map: twelve direct blocks, then one and two levels
        // of indirection.
        let per = self.block / 4;
        let ptr = |i: usize| u32_at(&inode.block, i * 4) as u64;
        let mut indirect = |table: u64, i: u64| -> Option<u64> {
            let mut b = [0u8; 4];
            read(table * self.block + i * 4, &mut b).ok()?;
            Some(u32::from_le_bytes(b) as u64).filter(|&p| p != 0)
        };
        let p = if n < 12 {
            ptr(n as usize)
        } else if n < 12 + per {
            indirect(ptr(12), n - 12)?
        } else if n < 12 + per + per * per {
            let k = n - 12 - per;
            let mid = indirect(ptr(13), k / per)?;
            indirect(mid, k % per)?
        } else {
            return None;
        };
        (p != 0).then_some(p)
    }

    /// Fill `buf` from the start of `inode`'s data.
    fn read_file(&self, read: &mut ReadFn, inode: &Inode, buf: &mut [u8]) -> Option<()> {
        if inode.flags & INLINE_DATA_FL != 0 || (buf.len() as u64) > inode.size.min(self.block) {
            return None;
        }
        let at = self.map(read, inode, 0)?;
        read(at * self.block, buf).ok()
    }

    /// Where a symlink points.
    fn link_target(&self, read: &mut ReadFn, inode: &Inode) -> Option<String> {
        if inode.size == 0 || inode.size > 4096 {
            return None;
        }
        let bytes = if inode.flags & (EXTENTS_FL | INLINE_DATA_FL) == 0 && inode.size < 60 {
            inode.block[..inode.size as usize].to_vec()
        } else {
            let mut b = vec![0u8; inode.size as usize];
            self.read_file(read, inode, &mut b)?;
            b
        };
        String::from_utf8(bytes).ok()
    }

    /// The inode named `name` in directory `dir`.
    fn lookup(&self, read: &mut ReadFn, dir: &Inode, name: &str) -> Option<u32> {
        if dir.mode & 0xF000 != 0x4000 || dir.flags & INLINE_DATA_FL != 0 || dir.size > MAX_DIR {
            return None;
        }
        let mut buf = vec![0u8; self.block as usize];
        for n in 0..dir.size.div_ceil(self.block) {
            let Some(at) = self.map(read, dir, n) else { continue };
            read(at * self.block, &mut buf).ok()?;
            let mut off = 0usize;
            while off + 8 <= buf.len() {
                let (ino, rec, len) = (u32_at(&buf, off), u16_at(&buf, off + 4) as usize, buf[off + 6] as usize);
                if rec < 8 || off + rec > buf.len() {
                    break;
                }
                if ino != 0 && off + 8 + len <= buf.len() && &buf[off + 8..off + 8 + len] == name.as_bytes() {
                    return Some(ino);
                }
                off += rec;
            }
        }
        None
    }

    /// The inode at `path` from the root, following every symlink.
    fn resolve(&self, read: &mut ReadFn, path: &str) -> Option<Inode> {
        let mut todo: Vec<String> = path.split('/').filter(|p| !p.is_empty()).rev().map(String::from).collect();
        let mut dirs = vec![ROOT_INODE];
        let mut links = 0;
        while let Some(name) = todo.pop() {
            match name.as_str() {
                "." => continue,
                ".." => {
                    if dirs.len() > 1 {
                        dirs.pop();
                    }
                    continue;
                }
                _ => {}
            }
            let dir = self.inode(read, *dirs.last()?)?;
            let n = self.lookup(read, &dir, &name)?;
            let inode = self.inode(read, n)?;
            if inode.mode & 0xF000 == 0xA000 {
                links += 1;
                if links > MAX_LINKS {
                    return None;
                }
                let target = self.link_target(read, &inode)?;
                if target.starts_with('/') {
                    dirs.truncate(1);
                }
                todo.extend(target.split('/').filter(|p| !p.is_empty()).rev().map(String::from));
                continue;
            }
            dirs.push(n);
        }
        self.inode(read, *dirs.last()?)
    }
}

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

    /// The first 64 bytes of an ELF program for `machine`, 64-bit little
    /// endian.
    fn elf(machine: u16) -> Vec<u8> {
        let mut h = vec![0u8; 64];
        h[..4].copy_from_slice(b"\x7fELF");
        h[4] = 2;
        h[5] = 1;
        h[6] = 1;
        h[16] = 3;
        h[18..20].copy_from_slice(&machine.to_le_bytes());
        h
    }

    /// An ext4 image of `tree`, as `flow rootfs import` makes one; `None`
    /// where mkfs.ext4 is missing.
    fn image(tree: &Path, out: &Path, size: &str, extra: &[&str]) -> Option<()> {
        let status = Command::new("mkfs.ext4")
            .args(["-q", "-F", "-L", "rootfs", "-E", "root_owner=0:0"])
            .args(extra)
            .arg("-d")
            .arg(tree)
            .arg(out)
            .arg(size)
            .status()
            .ok()?;
        assert!(status.success(), "mkfs.ext4 failed");
        Some(())
    }

    #[test]
    fn elf_headers_name_their_architecture() {
        assert_eq!(elf_arch(&elf(0x3E)), Some("x86_64"));
        assert_eq!(elf_arch(&elf(0xF3)), Some("riscv64"));
        assert_eq!(elf_arch(&elf(0xB7)), Some("aarch64"));
        assert_eq!(elf_arch(b"#!/bin/sh\necho hello\n"), None);
        assert_eq!(elf_arch(&elf(0x9999)), None);
    }

    #[test]
    fn an_images_shell_names_its_architecture_through_symlinks() {
        let dir = tempfile::tempdir().unwrap();
        // A merged /usr, as Debian and Ubuntu have it: bin -> usr/bin, and
        // sh -> dash; a riscv64 shell.
        let tree = dir.path().join("tree");
        std::fs::create_dir_all(tree.join("usr/bin")).unwrap();
        std::fs::write(tree.join("usr/bin/dash"), elf(0xF3)).unwrap();
        std::os::unix::fs::symlink("dash", tree.join("usr/bin/sh")).unwrap();
        std::os::unix::fs::symlink("usr/bin", tree.join("bin")).unwrap();
        // A full directory, so lookups cross blocks.
        for i in 0..400 {
            std::fs::write(tree.join(format!("usr/bin/tool-with-a-long-name-{i:04}")), b"#!/bin/sh\n").unwrap();
        }
        let img = dir.path().join("rootfs.ext4");
        if image(&tree, &img, "8M", &[]).is_none() {
            eprintln!("mkfs.ext4 not found: skipped");
            return;
        }
        assert_eq!(arch_of_file(&img), Some("riscv64"));
        // 1 KiB blocks, and an absolute link through a deeper directory.
        let tree2 = dir.path().join("tree2");
        std::fs::create_dir_all(tree2.join("bin")).unwrap();
        std::fs::create_dir_all(tree2.join("opt/busybox/bin")).unwrap();
        std::fs::write(tree2.join("opt/busybox/bin/busybox"), elf(0x3E)).unwrap();
        std::os::unix::fs::symlink("/opt/busybox/bin/../bin/busybox", tree2.join("bin/busybox")).unwrap();
        let img2 = dir.path().join("small.ext4");
        image(&tree2, &img2, "8M", &["-b", "1024"]).unwrap();
        assert_eq!(arch_of_file(&img2), Some("x86_64"), "no sh: busybox decides");
    }

    #[test]
    fn what_is_not_an_image_or_has_no_shell_answers_nothing() {
        let dir = tempfile::tempdir().unwrap();
        let junk = dir.path().join("junk");
        std::fs::write(&junk, vec![0x5Au8; 64 << 10]).unwrap();
        assert_eq!(arch_of_file(&junk), None);
        assert_eq!(arch_of_file(&dir.path().join("missing")), None);
        let tree = dir.path().join("tree");
        std::fs::create_dir_all(tree.join("bin")).unwrap();
        std::fs::write(tree.join("bin/sh"), b"#!/not/an/elf\n").unwrap();
        std::os::unix::fs::symlink("loop", tree.join("bin/loop")).unwrap();
        std::os::unix::fs::symlink("loop", tree.join("bin/busybox")).unwrap();
        let img = dir.path().join("rootfs.ext4");
        if image(&tree, &img, "8M", &[]).is_none() {
            return;
        }
        assert_eq!(arch_of_file(&img), None, "a script shell and a symlink loop name nothing");
    }
}