#!/usr/bin/env python3
"""012: Graphene hardened_malloc on LOS 23.2, the AXP GRAPHENE_MALLOC set.

Source of the integration is GrapheneOS platform_bionic branch 16:
64-bit libc uses hardened_malloc, 32-bit stays Scudo, and SurfaceFlinger
plus camera executables stay on Scudo. AXP documents that camera exception
as applying to every camera executable, which on sunfish is
/vendor/bin/hw/android.hardware.camera.provider@2.7-service-google rather
than the newer Pixel apex path.

Library pin: GrapheneOS/hardened_malloc 1d7fc7ffe0f8 (branch 14, 2025-04-03),
the newest branch that is not the Android 17 line. device_has_arm_mte is
removed because LOS soong has no such product variable and sunfish has no MTE.

Class region size is 2GiB, as DivestOS/AXP set it, because the sunfish
kernel has a 39-bit address space (GOLDEN: VA_BITS=39 with hardened_malloc).

system/core init.rc raises vm.max_map_count to 1048576, same as GrapheneOS 16.
frameworks/av no longer sets RLIMIT_AS on this tree, so that AXP/Pie hunk
does not apply.
"""
from __future__ import annotations

import os
import pathlib
import subprocess
import sys

MARKER = "SUNFISH_A16_012_HARDENED_MALLOC"
PIN = "1d7fc7ffe0f8068a32ab2ba33df7515e7bad3f5f"
REPO = "https://github.com/GrapheneOS/hardened_malloc.git"
REGION_48BIT = "-DCONFIG_CLASS_REGION_SIZE=34359738368"
REGION_39BIT = "-DCONFIG_CLASS_REGION_SIZE=2147483648"

HERE = pathlib.Path(__file__).resolve().parent
TREE = pathlib.Path(os.environ.get("ANDROID_BUILD_TOP", "/srv/rom/src/lineage"))


def must_replace(path: pathlib.Path, old: str, new: str, label: str) -> None:
    text = path.read_text()
    if new in text:
        print(f"already applied: {path} ({label})")
        return
    if old not in text:
        sys.exit(f"anchor missing for {label} in {path}")
    path.write_text(text.replace(old, new, 1))
    print(f"patched {label}: {path}")


def rewrite_region_comment(bp: pathlib.Path) -> None:
    text = bp.read_text()
    stale = f'{REGION_39BIT}", // 32GiB'
    fresh = f'{REGION_39BIT}", // 2GiB, 39-bit VA (DivestOS/AXP)'
    if stale not in text:
        return
    bp.write_text(text.replace(stale, fresh))
    print(f"region comment: {bp}")


def ensure_library() -> None:
    dest = TREE / "external/hardened_malloc"
    if (dest / ".git").is_dir():
        head = subprocess.check_output(
            ["git", "-C", dest, "rev-parse", "HEAD"], text=True
        ).strip()
        bp_text = (dest / "Android.bp").read_text()
        if head == PIN and MARKER in bp_text and REGION_39BIT in bp_text:
            rewrite_region_comment(dest / "Android.bp")
            print(f"hardened_malloc already at {PIN}")
            return
        subprocess.check_call(["git", "-C", dest, "fetch", "--depth", "1", "origin", PIN])
        subprocess.check_call(["git", "-C", dest, "checkout", "--detach", PIN])
    else:
        if dest.exists():
            sys.exit(f"{dest} exists and is not a git checkout")
        dest.parent.mkdir(parents=True, exist_ok=True)
        subprocess.check_call(
            ["git", "clone", "--depth", "1", REPO, str(dest)]
        )
        subprocess.check_call(["git", "-C", dest, "fetch", "--depth", "1", "origin", PIN])
        subprocess.check_call(["git", "-C", dest, "checkout", "--detach", PIN])
    bp = dest / "Android.bp"
    text = bp.read_text()
    mte = """        device_has_arm_mte: {
            cflags: ["-DHAS_ARM_MTE", "-march=armv8-a+dotprod+memtag"]
        },
"""
    if mte in text:
        text = text.replace(mte, "")
    # sunfish msm-4.14 is CONFIG_ARM64_VA_BITS=39. The 32GiB default class
    # region needs 48-bit VA and aborts every 64-bit process at boot.
    # DivestOS/AXP Patch.sh: sed 's/34359738368/2147483648/' Android.bp.
    text = text.replace(REGION_48BIT, REGION_39BIT)
    text = text.replace(
        f'{REGION_39BIT}", // 32GiB',
        f'{REGION_39BIT}", // 2GiB, 39-bit VA (DivestOS/AXP)',
    )
    bp.write_text(text)
    text = bp.read_text()
    if REGION_39BIT not in text:
        sys.exit(f"CONFIG_CLASS_REGION_SIZE is not 2GiB in {bp}")
    if MARKER not in text:
        text = text.replace(
            "cc_library {\n",
            f"cc_library {{\n    // {MARKER}: sunfish has no MTE; LOS soong has no device_has_arm_mte.\n",
            1,
        )
    bp.write_text(text)
    head = subprocess.check_output(
        ["git", "-C", dest, "rev-parse", "HEAD"], text=True
    ).strip()
    if head != PIN:
        sys.exit(f"hardened_malloc HEAD {head} != {PIN}")
    print(f"hardened_malloc {PIN}")


def patch_bionic() -> None:
    wrapper = TREE / "bionic/libc/bionic/h_malloc_wrapper.cpp"
    src = HERE / "h_malloc_wrapper.cpp"
    if not wrapper.exists() or wrapper.read_text() != src.read_text():
        wrapper.write_text(src.read_text())
        print(f"installed {wrapper}")

    must_replace(
        TREE / "bionic/libc/Android.bp",
        """cc_defaults {
    name: "libc_native_allocator_defaults",

    whole_static_libs: [
        "libscudo",
    ],
    cflags: [
        "-DUSE_SCUDO",
    ],
    header_libs: ["gwp_asan_headers"],
    product_variables: {
        malloc_low_memory: {
            cflags: ["-UUSE_SCUDO"],
            whole_static_libs: [
                "libjemalloc5",
                "libc_jemalloc_wrapper",
            ],
            exclude_static_libs: [
                "libscudo",
            ],
        },
    },
}
""",
        f"""cc_defaults {{
    // {MARKER}
    name: "libc_native_allocator_defaults",

    header_libs: ["gwp_asan_headers"],
    whole_static_libs: ["libscudo"],

    // GrapheneOS 16: hardened_malloc on 64-bit, Scudo on 32-bit.
    multilib: {{
        lib32: {{
            cflags: ["-DUSE_SCUDO"],
        }},
        lib64: {{
            srcs: ["bionic/h_malloc_wrapper.cpp"],
            cflags: [
                "-DH_MALLOC_PREFIX",
                "-DUSE_H_MALLOC",
                "-DUSE_SCUDO",
            ],
            whole_static_libs: [
                "libhardened_malloc",
            ],
        }},
    }},
}}
""",
        "allocator-defaults",
    )

    must_replace(
        TREE / "bionic/libc/bionic/malloc_common.h",
        """#else // __has_feature(hwaddress_sanitizer)

#if defined(USE_SCUDO)
""",
        f"""#else // __has_feature(hwaddress_sanitizer)

// {MARKER}
#ifdef __LP64__
#ifndef USE_H_MALLOC
#error missing USE_H_MALLOC
#endif

#include "h_malloc.h"
#define Malloc(function)  h_ ## function
__BEGIN_DECLS
int h_malloc_info(int options, FILE* fp);
__END_DECLS

#if defined(USE_SCUDO)
#include "scudo.h"
void InitNativeAllocatorDispatch(libc_globals* globals);
#endif

#define BOTH_H_MALLOC_AND_SCUDO

#else // 32-bit

#if defined(USE_SCUDO)
""",
        "malloc-header",
    )
    header = TREE / "bionic/libc/bionic/malloc_common.h"
    text = header.read_text()
    tail = """#define Malloc(function)  je_ ## function

#endif

#endif

const MallocDispatch* NativeAllocatorDispatch();
"""
    tail_new = """#define Malloc(function)  je_ ## function

#endif

#endif // __LP64__

#endif // __has_feature(hwaddress_sanitizer)

const MallocDispatch* NativeAllocatorDispatch();
"""
    if tail_new not in text:
        if tail not in text:
            sys.exit(f"header close anchor missing in {header}")
        header.write_text(text.replace(tail, tail_new, 1))

    must_replace(
        TREE / "bionic/libc/bionic/malloc_common.cpp",
        """const MallocDispatch* NativeAllocatorDispatch() {
  return &__libc_malloc_default_dispatch;
}
""",
        f"""#if defined(BOTH_H_MALLOC_AND_SCUDO)
// {MARKER}: camera executables and surfaceflinger stay on Scudo.

#define ScudoMalloc(function)  scudo_ ## function

static constexpr MallocDispatch __scudo_malloc_dispatch __attribute__((unused)) = {{
  ScudoMalloc(calloc),
  ScudoMalloc(free),
  ScudoMalloc(mallinfo),
  ScudoMalloc(malloc),
  ScudoMalloc(malloc_usable_size),
  ScudoMalloc(memalign),
  ScudoMalloc(posix_memalign),
#if defined(HAVE_DEPRECATED_MALLOC_FUNCS)
  ScudoMalloc(pvalloc),
#endif
  ScudoMalloc(realloc),
#if defined(HAVE_DEPRECATED_MALLOC_FUNCS)
  ScudoMalloc(valloc),
#endif
  ScudoMalloc(malloc_iterate),
  ScudoMalloc(malloc_disable),
  ScudoMalloc(malloc_enable),
  ScudoMalloc(mallopt),
  ScudoMalloc(aligned_alloc),
  ScudoMalloc(malloc_info),
}};

static const MallocDispatch* native_allocator_dispatch;

void InitNativeAllocatorDispatch(libc_globals* globals) {{
  bool hardened_impl = true;
  if (globals->flags & GLOBAL_FLAG_DISABLE_HARDENED_MALLOC) {{
    hardened_impl = false;
  }} else if (getenv("DISABLE_HARDENED_MALLOC") != nullptr) {{
    hardened_impl = false;
  }}

  const MallocDispatch* table = hardened_impl ?
    &__libc_malloc_default_dispatch :
    &__scudo_malloc_dispatch;

  if (!hardened_impl) {{
    globals->malloc_dispatch_table = __scudo_malloc_dispatch;
    globals->current_dispatch_table = &globals->malloc_dispatch_table;
    globals->default_dispatch_table = &globals->malloc_dispatch_table;
  }}

  native_allocator_dispatch = table;
}}

const MallocDispatch* NativeAllocatorDispatch() {{
  return native_allocator_dispatch;
}}
#else
const MallocDispatch* NativeAllocatorDispatch() {{
  return &__libc_malloc_default_dispatch;
}}
#endif
""",
        "dispatch",
    )

    must_replace(
        TREE / "bionic/libc/bionic/malloc_common_dynamic.cpp",
        """void MallocInitImpl(libc_globals* globals) {
  char prop[PROP_VALUE_MAX];
""",
        f"""void MallocInitImpl(libc_globals* globals) {{
#if defined(BOTH_H_MALLOC_AND_SCUDO)
  // {MARKER}
  InitNativeAllocatorDispatch(globals);
#endif
  char prop[PROP_VALUE_MAX];
""",
        "malloc-init",
    )

    must_replace(
        TREE / "bionic/libc/private/bionic_globals.h",
        """  MallocDispatch malloc_dispatch_table;
};
""",
        f"""  MallocDispatch malloc_dispatch_table;
  // {MARKER}
  int flags;
}};

#define GLOBAL_FLAG_DISABLE_HARDENED_MALLOC 1
""",
        "globals",
    )

    init = TREE / "bionic/libc/bionic/libc_init_dynamic.cpp"
    text = init.read_text()
    if '#include <string.h>' not in text:
        if '#include <stdlib.h>' not in text:
            sys.exit(f"stdlib include missing in {init}")
        text = text.replace(
            "#include <stdlib.h>\n",
            "#include <stdlib.h>\n#include <string.h>\n#include <unistd.h>\n",
            1,
        )
    helper = f"""
// {MARKER}: GrapheneOS disables hardened_malloc for surfaceflinger.
// AXP expands the camera exception to every camera executable. sunfish runs
// /vendor/bin/hw/android.hardware.camera.provider@2.7-service-google.
static void init_hardened_malloc_exclusions(libc_globals* globals) {{
  char exe_path[500];
  ssize_t n = readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1);
  if (n <= 0) {{
    globals->flags = 0;
    return;
  }}
  exe_path[n] = '\\0';
  if (strcmp(exe_path, "/system/bin/surfaceflinger") == 0 ||
      strstr(exe_path, "camera") != nullptr) {{
    globals->flags = GLOBAL_FLAG_DISABLE_HARDENED_MALLOC;
  }} else {{
    globals->flags = 0;
  }}
}}

"""
    # __libc_preinit_impl must stay noinline (stack protector guard setup).
    noinline = "__attribute__((noinline))\n"
    anchor = noinline + "static void __libc_preinit_impl() {"
    misplaced = noinline + helper + "static void __libc_preinit_impl() {"
    if misplaced in text:
        text = text.replace(misplaced, helper + anchor, 1)
    if "init_hardened_malloc_exclusions" not in text:
        if anchor not in text:
            sys.exit(f"preinit anchor missing in {init}")
        text = text.replace(anchor, helper + anchor, 1)
    if anchor not in text:
        sys.exit(f"__libc_preinit_impl lost noinline in {init}")
    old_mutate = "  __libc_globals.mutate(__libc_init_malloc);\n"
    new_mutate = """  __libc_globals.mutate([](libc_globals* globals) {
    init_hardened_malloc_exclusions(globals);
    __libc_init_malloc(globals);
  });
"""
    if "init_hardened_malloc_exclusions(globals)" not in text:
        if old_mutate not in text:
            sys.exit(f"mutate anchor missing in {init}")
        text = text.replace(old_mutate, new_mutate, 1)
    init.write_text(text)
    print(f"patched exclusions: {init}")


def patch_init_rc() -> None:
    must_replace(
        TREE / "system/core/rootdir/init.rc",
        "    symlink /sys/class/block/${dev.mnt.rootdisk.data} /dev/sys/block/by-name/rootdisk\n",
        "    symlink /sys/class/block/${dev.mnt.rootdisk.data} /dev/sys/block/by-name/rootdisk\n\n"
        f"    # {MARKER}: GrapheneOS 16 vm.max_map_count for hardened_malloc\n"
        "    write /proc/sys/vm/max_map_count 1048576\n",
        "max-map-count",
    )


def main() -> None:
    if not (TREE / "bionic/libc/Android.bp").is_file():
        sys.exit(f"not an Android tree: {TREE}")
    ensure_library()
    patch_bionic()
    patch_init_rc()
    print(f"done: {MARKER}")


if __name__ == "__main__":
    main()