#!/usr/bin/bash
#
# Description: fsck.ubifs (mtd-utils 2.3.1) identifies corruptions in a newly encrypted UBIFS
# volume created with mkfs.ubifs (mtd-utils 2.1.1).
#
# Script Functionality:
#   1. Builds mkfs.ubifs/ubiformat/ubiattach etc.from mtd-utils v2.1.1 
#      and mtd-utils v2.3.1 plus fsck.ubifs from mtd-utils v2.3.1.
#   2. Creates an emulated MTD device (mtdram) and an UBI volume on top of it.
#   3. format and encrypt the UBI volume directly from that directory with the OLD mkfs.ubifs. 
#   4. Mounts it read-only as a sanity check, then unmounts.
#   5. Runs the NEW fsck.ubifs against the volume, capturing whether it
#      flags corruption / attempts a rebuild.
#
# Usage: fsck_ubifs_reproduce_bug.sh [--encrypted|--unencrypted] [--apply-patch|--no-patch] [path]

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PATCH_FILE="$SCRIPT_DIR/0001-fsck.ubifs-fix-for-pre-mkfs.ubifs-2.2.1-xattr-dentry-type.patch"

# Command-line: whether to encrypt the ubi volume (default) or leave it unencrypted (for comparison).
DO_ENCRYPT=1
# Command-line: whether to apply the fsck.ubifs patch to the NEW (2.3.1) build before building it.
APPLY_PATCH=0
WORKDIR=""
usage() {
    echo "Usage: $0 [--encrypted|--unencrypted] [--apply-patch|--no-patch] [path]"
    echo "  --encrypted    format the volume with mkfs.ubifs encryption (default)"
    echo "  --unencrypted  format the volume with no encryption at all"
    echo "  --apply-patch  apply $(basename "$PATCH_FILE") to mtd-utils v2.3.1 before building it"
    echo "  --no-patch     build mtd-utils v2.3.1 unpatched (default)"
    echo "  path           The absolute path to the directory where the test will be staged"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --encrypted)
            DO_ENCRYPT=1
            shift ;;
        --unencrypted)
            DO_ENCRYPT=0
            shift ;;
        --apply-patch)
            APPLY_PATCH=1
            shift ;;
        --no-patch)
            APPLY_PATCH=0
            shift ;;
        -h|--help)
            usage; exit 0 ;;
        -*)
            echo "Unknown option: $1" >&2; usage; exit 1 ;;
        *)
            WORKDIR="$1"
            shift ;;
    esac
done

if [[ -z "$WORKDIR" ]]; then
    echo "Error: path is required" >&2
    usage
    exit 1
fi

if [[ "$APPLY_PATCH" -eq 1 && ! -f "$PATCH_FILE" ]]; then
    echo "Error: patch file not found at $PATCH_FILE" >&2
    exit 1
fi

# Parameters for the test, these are mostly arbitrary aside from the mtd-utils versions 
MTD_UTILS_GIT="${MTD_UTILS_GIT:-git://git.infradead.org/mtd-utils.git}"
MTD_UTILS_OLD_TAG="v2.1.1"
MTD_UTILS_NEW_TAG="v2.3.1"
MTD_SIZE_KIB=$((32 * 1024))   
ERASE_SIZE_KIB=64             
MOUNTPOINT="$WORKDIR/mnt"

MTD_UTILS_OLD_PREFIX="$WORKDIR/mtd-utils-old"
if [[ "$APPLY_PATCH" -eq 1 ]]; then
    MTD_UTILS_NEW_PREFIX="$WORKDIR/mtd-utils-new-patched"
else
    MTD_UTILS_NEW_PREFIX="$WORKDIR/mtd-utils-new-unpatched"
fi

mkdir -p "$WORKDIR" "$MOUNTPOINT"

# Clean the runtime environment on exit
cleanup()
{
    echo "cleaning up"
    set +e

    umount "$MOUNTPOINT" 2>/dev/null
    rmdir "$MOUNTPOINT" 2>/dev/null

    if [[ -n "${UBI_NUM:-}" ]]; then
        echo "detaching UBI device ubi$UBI_NUM"
        "$UBIDETACH" -d "$UBI_NUM"
    fi

    echo "removing ubi module"
    modprobe -r ubi 
    echo "removing mtdram module"
    modprobe -r mtdram
}
trap cleanup EXIT

# Build OLD (2.1.1) and NEW (2.3.1) mtd-utils
build_mtd_utils()
{
    local tag="$1" prefix="$2" patch="${3:-}"
    local srcDir="$WORKDIR/src-$(basename "$prefix")"
    if [[ -x "$prefix/sbin/mkfs.ubifs" && -x "$prefix/sbin/fsck.ubifs" ]]; then
        echo "mtd-utils $tag already built at $prefix, skipping"
        return
    fi
    echo "cloning/building mtd-utils $tag into $prefix"
    if [[ ! -d "$srcDir" ]]; then
        git clone --branch "$tag" --depth 1 "$MTD_UTILS_GIT" "$srcDir"
        if [[ -n "$patch" ]]; then
            echo "applying patch $(basename "$patch") to $srcDir"
            ( cd "$srcDir" && patch -p1 < "$patch" )
        fi
        ( cd "$srcDir" && \
        ./autogen.sh && \
        ./configure --prefix="$prefix" --without-xattr && \
        make -j"$(nproc)"  && \
        make install  )
    fi
    echo "built mtd-utils $tag -> $prefix"
}

build_mtd_utils "$MTD_UTILS_OLD_TAG" "$MTD_UTILS_OLD_PREFIX"
if [[ "$APPLY_PATCH" -eq 1 ]]; then
    build_mtd_utils "$MTD_UTILS_NEW_TAG" "$MTD_UTILS_NEW_PREFIX" "$PATCH_FILE"
else
    build_mtd_utils "$MTD_UTILS_NEW_TAG" "$MTD_UTILS_NEW_PREFIX"
fi

MKFS_OLD="$MTD_UTILS_OLD_PREFIX/sbin/mkfs.ubifs"
UBIATTACH="$MTD_UTILS_NEW_PREFIX/sbin/ubiattach"       
UBIDETACH="$MTD_UTILS_NEW_PREFIX/sbin/ubidetach"
UBIMKVOL="$MTD_UTILS_NEW_PREFIX/sbin/ubimkvol"
FSCK_NEW="$MTD_UTILS_NEW_PREFIX/sbin/fsck.ubifs"

echo "OLD mkfs.ubifs: $($MKFS_OLD --version 2>&1 | head -1)"
echo "NEW fsck.ubifs: $($FSCK_NEW --version 2>&1 | head -1)"

# setup mtdram and the ubi volume
setup_mtd() {
    if modprobe mtdram total_size="$MTD_SIZE_KIB" erase_size="$ERASE_SIZE_KIB"; then
        echo "mtdram loaded"
    fi

    sleep 1
    MTD=$(cat /proc/mtd | grep "mtdram" | tr ":" " " | awk '{ print $1 }')
    MTD_NUM=$(echo "$MTD" | tr -dc '0-9')
    
    if [[ -z "${MTD_NUM:-}" ]]; then
        echo "could not find emulated mtd device" >&2
        exit 1
    fi

    MTD_NODE="/dev/$MTD"
    echo "using $MTD_NODE"
}

setup_ubi() {
    local attach_out
    if modprobe ubi; then
        echo "ubi module loaded"
    fi

    sleep 1
    attach_out=$("$UBIATTACH" -p "$MTD_NODE") || true
    echo "$attach_out"
    UBI_NUM=$(echo "$attach_out" | tr "," " " | awk '{ print $4 }')
    
    if [[ -z "${UBI_NUM:-}" ]]; then
        echo "could not attach $MTD_NODE to UBI" >&2
        exit 1
    fi

    sleep 1
    "$UBIMKVOL" "/dev/ubi$UBI_NUM" -N test_vol -m 
    sleep 1
    VOL_DEV=$(ls /dev/ubi${UBI_NUM}_* | head -1)
    echo "UBI volume device: $VOL_DEV"
}

setup_mtd
setup_ubi

# Stage test files, then format+populate the UBI volume, and encrypt it with the OLD mkfs.ubifs
STAGING_DIR="$WORKDIR/ubifs_root"
rm -rf "$STAGING_DIR"
mkdir -p "$STAGING_DIR"

MKFS_UBIFS_PARAMS="-m 2048 -e 129024 -x zlib"
MAX_LEB=2047

ENCRYPT_ARGS=()
if [[ "$DO_ENCRYPT" -eq 1 ]]; then
    # Dummy raw key file for test purposes. Use random bytes to feed mkfs.ubifs's -K option. mkfs.ubifs requires at least 64 bytes.
    KEYFILE="$WORKDIR/raw.key"
    head -c 64 /dev/urandom > "$KEYFILE"
    chmod 600 "$KEYFILE"
    ENCRYPT_ARGS=(-K "$KEYFILE" -C AES-256-XTS)
    echo "formatting+encrypting $VOL_DEV with OLD mkfs.ubifs ($MTD_UTILS_OLD_TAG) from $STAGING_DIR"
else
    echo "formatting (unencrypted) $VOL_DEV with OLD mkfs.ubifs ($MTD_UTILS_OLD_TAG) from $STAGING_DIR"
fi

"$MKFS_OLD" $MKFS_UBIFS_PARAMS -r "$STAGING_DIR" "${ENCRYPT_ARGS[@]}" -o "$VOL_DEV" -c "$MAX_LEB"

# Mount read-only as a sanity check
echo "mounting $VOL_DEV at $MOUNTPOINT (read-only sanity check)"
mount -t ubifs -o ro "$VOL_DEV" "$MOUNTPOINT"

echo "directory listing under encryption (kernel-decrypted view):"
ls -la "$MOUNTPOINT" 

# Clean unmount, then run the NEW fsck.ubifs
echo "unmounting cleanly"
umount "$MOUNTPOINT"

echo "running NEW fsck.ubifs ($MTD_UTILS_NEW_TAG) against $VOL_DEV"
set +e
"$FSCK_NEW" -y "$VOL_DEV" > "$WORKDIR/fsck_output.txt" 2>&1
FSCK_RC=$?
set -e
echo "fsck.ubifs exit code: $FSCK_RC"
echo "fsck.ubifs output saved to $WORKDIR/fsck_output.txt"
cat "$WORKDIR/fsck_output.txt"