"""
Test lldb's support for the AArch64 Permission Overlay extension (POE), which
is used to implement Linux's memory protection keys feature.
"""

import lldb
from lldbsuite.test.decorators import *
from lldbsuite.test.lldbtest import *
from lldbsuite.test import lldbutil


class AArch64LinuxPOE(TestBase):
    NO_DEBUG_INFO_TESTCASE = True

    EXPECTED_POR_EL0 = "por_el0 = 0x0000000001234567"
    EXPECTED_POR_EL0_FIELDS = (
        "          = {\n"
        "              Perm15 = No Access\n"
        "              Perm14 = No Access\n"
        "              Perm13 = No Access\n"
        "              Perm12 = No Access\n"
        "              Perm11 = No Access\n"
        "              Perm10 = No Access\n"
        "              Perm9 = No Access\n"
        "              Perm8 = No Access\n"
        "              Perm7 = No Access\n"
        "              Perm6 = Read\n"
        "              Perm5 = Execute\n"
        "              Perm4 = Read, Execute\n"
        "              Perm3 = Write\n"
        "              Perm2 = Write, Read\n"
        "              Perm1 = Write, Execute\n"
        "              Perm0 = Read, Write, Execute\n"
        "            }"
    )

    @skipUnlessArch("aarch64")
    @skipUnlessPlatform(["linux"])
    def test_poe_live(self):
        if not self.isAArch64POE():
            self.skipTest("POE must be present.")

        self.build()
        self.runCmd("file " + self.getBuildArtifact("a.out"), CURRENT_EXECUTABLE_SET)

        lldbutil.run_break_set_by_file_and_line(
            self,
            "main.c",
            line_number("main.c", "// Set break point at this line."),
            num_expected_locations=1,
        )

        self.runCmd("run", RUN_SUCCEEDED)

        if self.process().GetState() == lldb.eStateExited:
            self.fail("Test program failed to run.")

        self.expect(
            "thread list",
            STOPPED_DUE_TO_BREAKPOINT,
            substrs=["stopped", "stop reason = breakpoint"],
        )

        self.expect(
            "register read --all",
            substrs=[
                "Permission Overlay Registers",
                f"{self.EXPECTED_POR_EL0}",
            ],
        )

        if self.hasXMLSupport():
            self.expect(
                "register read por_el0",
                substrs=[f" {self.EXPECTED_POR_EL0}\n" + self.EXPECTED_POR_EL0_FIELDS],
            )

        # POR_EL0 should be restored after expression evaluation.
        self.expect("expression expr_function()", substrs=["$0 = 1"])
        self.expect("register read por_el0", substrs=[self.EXPECTED_POR_EL0])

        # Unmapped region has no key (not even default).
        self.expect("memory region 0", substrs=["protection key:"], matching=False)

        # The region has base permissions r-x, and overlay is r--. The result
        # is that execution is disabled.
        self.expect(
            "memory region read_only_page",
            substrs=["rw-", "protection key: 6 (r--, effective: r--)"],
        )
        # A region not assigned to a protection key has the default key 0. This
        # key is rwx, but overlays cannot add permissions not already in the
        # page table. So the execute permission is not enabled.
        self.expect(
            "memory region key_zero_page",
            substrs=["rw-", "protection key: 0 (rwx, effective: rw-)"],
        )

        # Overlay permissions are on their own line.
        self.expect(
            "memory region --all",
            patterns=["\nprotection key: [0-9]+ \([rwx-]{3}, effective: [rwx-]{3}\)\n"],
        )

        # Not passing this to the application allows us to fix the permissions
        # using lldb, then continue to a normal exit.
        self.runCmd("process handle SIGSEGV --pass false")

        self.expect(
            "continue",
            substrs=[
                "stop reason = signal SIGSEGV: failed protection key checks (fault address="
            ],
        )

        # This fault should have happened due to the write, not the read.
        self.assertEqual(
            self.dbg.GetSelectedTarget()
            .GetProcess()
            .GetSelectedThread()
            .GetSelectedFrame()
            .GetFunctionName(),
            "cause_write_fault",
        )

        # Allow writes so we can continue. This value has permission 6 changed
        # from read only (1) to write (4).
        self.runCmd("register write por_el0 0x4234567")

        self.expect("continue", substrs=["exited with status = 0"])

    @skipIfLLVMTargetMissing("AArch64")
    def test_poe_core(self):
        # Core was generated by running the test program on a system with POE
        # and coredump_filter set to 0.
        self.runCmd("target create --core corefile")

        self.expect("process status", substrs=["SIGSEGV: failed protection key checks"])

        self.expect(
            "register read --all",
            substrs=[
                "Permission Overlay Registers",
                f"{self.EXPECTED_POR_EL0}",
            ],
        )

        if self.hasXMLSupport():
            self.expect(
                "register read por_el0",
                substrs=[f" {self.EXPECTED_POR_EL0}\n" + self.EXPECTED_POR_EL0_FIELDS],
            )

            # por_el0 is an unusal case where every field uses the same enum.
            # We should print all the field names in a list, then the enum only
            # once. Rather than printing the enum 15 times, once for each field.
            self.expect(
                "register info por_el0",
                substrs=[
                    ", ".join([f"Perm{n}" for n in range(15, -1, -1)])
                    + ": 0 = No Access"
                ],
            )

        # Protection keys are listed in /proc/<pid>/smaps, which is not included
        # in core files.
        self.expect("memory region --all", substrs=["protection key:"], matching=False)
