#!/bin/sh
set -eu

case "$(uname -s)" in
	Linux) ;;
	*)
		echo "PoC requires Linux LD_PRELOAD. Run this on Debian/Linux."
		exit 0
		;;
esac

parted_bin=${PARTED_BIN:-}
if [ -z "$parted_bin" ]; then
	if which parted >/dev/null 2>&1; then
		parted_bin=$(which parted)
	fi
fi

if [ -z "$parted_bin" ]; then
	echo "PoC skipped: set PARTED_BIN to a GNU parted binary, e.g. /usr/sbin/parted."
	exit 0
fi

cc_bin=${CC:-cc}
if ! which "$cc_bin" >/dev/null 2>&1; then
	echo "PoC requires a C compiler. Set CC=/path/to/cc if needed."
	exit 1
fi

if ! which python3 >/dev/null 2>&1; then
	echo "PoC requires python3."
	exit 1
fi

tmp=$(mktemp -d "${TMPDIR:-/tmp}/parted-ped-calloc-poc.XXXXXX")
trap 'rm -rf "$tmp"' EXIT
lib="$tmp/fail-ped-malloc.so"
img=${IMAGE_PATH:-./ped-calloc-null-mbr.img}

cat >"$tmp/fail-ped-malloc.c" <<'C'
#define _GNU_SOURCE
#include <dlfcn.h>
#include <errno.h>
#include <execinfo.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static void *(*real_ped_malloc)(size_t);
static int failed;

static int
stack_has_ped_calloc(void)
{
	void *frames[48];
	int n = backtrace(frames, 48);

	for (int i = 0; i < n; i++) {
		Dl_info info;
		if (dladdr(frames[i], &info) == 0 || info.dli_sname == NULL)
			continue;
		if (strstr(info.dli_sname, "ped_calloc") != NULL)
			return 1;
	}

	return 0;
}

void *
ped_malloc(size_t size)
{
	if (real_ped_malloc == NULL) {
		real_ped_malloc = dlsym(RTLD_NEXT, "ped_malloc");
		if (real_ped_malloc == NULL) {
			fprintf(stderr, "fail-ped-malloc: could not find real ped_malloc\n");
			_exit(127);
		}
	}

	if (!failed && stack_has_ped_calloc()) {
		failed = 1;
		errno = ENOMEM;
		fprintf(stderr,
		        "fail-ped-malloc: returning NULL for ped_malloc(%zu) from ped_calloc stack\n",
		        size);
		return NULL;
	}

	return real_ped_malloc(size);
}
C

"$cc_bin" -shared -fPIC "$tmp/fail-ped-malloc.c" -o "$lib" -ldl

python3 - "$img" <<'PY'
import struct
import sys

img = sys.argv[1]
sector_size = 512
sectors = 20480

with open(img, "wb") as f:
    f.truncate(sectors * sector_size)
    mbr = bytearray(sector_size)
    # One normal Linux partition. Parsing it reaches msdos_partition_new(),
    # which allocates DosPartitionData with ped_calloc().
    mbr[446:462] = struct.pack(
        "<B3sB3sII",
        0x00,
        b"\x00\x02\x00",
        0x83,
        b"\xff\xff\xff",
        2048,
        4096,
    )
    mbr[510:512] = b"\x55\xaa"
    f.seek(0)
    f.write(mbr)

print(f"created {img}")
print("image layout: raw MBR disk image with one Linux partition")
PY

echo
echo "Running with ped_malloc failure injected only from a ped_calloc() call stack:"
echo "  $parted_bin -s $img unit s print"

preload=$lib
if [ -n "${ASAN_PRELOAD:-}" ]; then
	preload="$ASAN_PRELOAD:$preload"
elif [ -n "${LD_PRELOAD:-}" ]; then
	preload="$LD_PRELOAD:$preload"
fi

set +e
LD_PRELOAD="$preload" \
ASAN_OPTIONS="${ASAN_OPTIONS:-abort_on_error=1:symbolize=1}" \
	"$parted_bin" -s "$img" unit s print >"$tmp/stdout" 2>"$tmp/stderr"
status=$?
set -e

echo "parted exit status: $status"
if [ -s "$tmp/stdout" ]; then
	echo "stdout:"
	sed 's/^/  /' "$tmp/stdout"
fi
if [ -s "$tmp/stderr" ]; then
	echo "stderr:"
	sed 's/^/  /' "$tmp/stderr"
fi

if ! grep -q 'fail-ped-malloc: returning NULL' "$tmp/stderr"; then
	echo "Bug not triggered: ped_malloc failure did not occur from a visible ped_calloc() stack."
	echo "Use an unstripped/shared GNU Parted build, preferably the ASAN build from ./asan-build-and-run.sh."
	exit 1
fi

case "$status" in
	139|134)
		echo "BUG TRIGGERED: real GNU Parted crashed after ped_calloc() received NULL."
		exit 0
		;;
	*)
		if grep -qi 'AddressSanitizer\|SEGV\|segmentation fault\|null pointer' "$tmp/stderr" 2>/dev/null; then
			echo "BUG TRIGGERED: sanitizer/crash text was reported."
			exit 0
		fi
		echo "Malloc failure was injected, but this binary did not crash in the expected way."
		exit 1
		;;
esac