//! 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");

    }

}