Fix potential stack corruption in DebugSessionImp

- allocate 64 bytes on stack when accessing registers

Signed-off-by: Mateusz Hoppe <[email protected]>

Author: HoppeMateusz

level_zero/tools/source/debug/debug_session_imp.cpp

@@ -343,14 +343,18 @@ bool DebugSessionImp::writeResumeCommand(const std::vector<EuThread::ThreadId> &
                 dword = 4;
             }
 
+            const auto regSize = std::max(getRegisterSize(registerType), hwInfo.capabilityTable.grfSize);
+            auto reg = std::make_unique<uint32_t[]>(regSize / sizeof(uint32_t));
+
             for (auto &threadID : threadIds) {
-                uint32_t reg[8] = {};
                 auto apiThread = convertToApi(threadID);
-                if (readRegistersImp(apiThread, registerType, 0, 1, reg) != ZE_RESULT_SUCCESS) {
+                memset(reg.get(), 0, regSize);
+
+                if (readRegistersImp(apiThread, registerType, 0, 1, reg.get()) != ZE_RESULT_SUCCESS) {
                     success = false;
                 } else {
                     reg[dword] |= sipResumeValue;
-                    if (writeRegistersImp(apiThread, registerType, 0, 1, reg) != ZE_RESULT_SUCCESS) {
+                    if (writeRegistersImp(apiThread, registerType, 0, 1, reg.get()) != ZE_RESULT_SUCCESS) {
                         success = false;
                     }
                 }
@@ -626,14 +630,16 @@ bool DebugSessionImp::isForceExceptionOrForceExternalHaltOnlyExceptionReason(uin
 
 void DebugSessionImp::fillResumeAndStoppedThreadsFromNewlyStopped(std::vector<EuThread::ThreadId> &resumeThreads, std::vector<EuThread::ThreadId> &stoppedThreadsToReport) {
 
+    const auto regSize = std::max(getRegisterSize(ZET_DEBUG_REGSET_TYPE_CR_INTEL_GPU), 64u);
+    auto reg = std::make_unique<uint32_t[]>(regSize / sizeof(uint32_t));
+
     for (auto &newlyStopped : newlyStoppedThreads) {
         if (allThreads[newlyStopped]->isStopped()) {
-            uint32_t reg[8] = {};
-
+            memset(reg.get(), 0, regSize);
             ze_device_thread_t apiThread = convertToApi(newlyStopped);
-            readRegistersImp(apiThread, ZET_DEBUG_REGSET_TYPE_CR_INTEL_GPU, 0, 1, reg);
+            readRegistersImp(apiThread, ZET_DEBUG_REGSET_TYPE_CR_INTEL_GPU, 0, 1, reg.get());
 
-            if (isForceExceptionOrForceExternalHaltOnlyExceptionReason(reg)) {
+            if (isForceExceptionOrForceExternalHaltOnlyExceptionReason(reg.get())) {
                 PRINT_DEBUGGER_THREAD_LOG("RESUME accidentally stopped thread = %s\n", allThreads[newlyStopped]->toString().c_str());
                 resumeThreads.push_back(newlyStopped);
             } else {
@@ -727,7 +733,10 @@ const SIP::regset_desc *DebugSessionImp::getSbaRegsetDesc() {
 const SIP::regset_desc *DebugSessionImp::typeToRegsetDesc(uint32_t type) {
     auto pStateSaveAreaHeader = getStateSaveAreaHeader();
 
-    UNRECOVERABLE_IF(pStateSaveAreaHeader == nullptr);
+    DEBUG_BREAK_IF(pStateSaveAreaHeader == nullptr);
+    if (pStateSaveAreaHeader == nullptr) {
+        return nullptr;
+    }
 
     switch (type) {
     case ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU:
@@ -757,6 +766,14 @@ const SIP::regset_desc *DebugSessionImp::typeToRegsetDesc(uint32_t type) {
     }
 }
 
+uint32_t DebugSessionImp::getRegisterSize(uint32_t type) {
+    auto regset = typeToRegsetDesc(type);
+    if (regset) {
+        return regset->bytes;
+    }
+    return 0;
+}
+
 uint32_t DebugSessionImp::typeToRegsetFlags(uint32_t type) {
     switch (type) {
     case ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU:
@@ -806,8 +823,11 @@ ze_result_t DebugSessionImp::readSbaRegisters(ze_device_thread_t thread, uint32_
         return ret;
     }
 
-    uint32_t r0[8];
-    ret = readRegistersImp(thread, ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU, 0, 1, r0);
+    const auto &hwInfo = connectedDevice->getHwInfo();
+    const auto regSize = std::max(getRegisterSize(ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU), hwInfo.capabilityTable.grfSize);
+    auto r0 = std::make_unique<uint32_t[]>(regSize / sizeof(uint32_t));
+
+    ret = readRegistersImp(thread, ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU, 0, 1, r0.get());
     if (ret != ZE_RESULT_SUCCESS) {
         return ret;
     }

level_zero/tools/source/debug/debug_session_imp.h

@@ -79,6 +79,8 @@ struct DebugSessionImp : DebugSession {
                                       uint32_t start, uint32_t count, void *pRegisterValues, bool write);
 
     const SIP::regset_desc *typeToRegsetDesc(uint32_t type);
+    uint32_t getRegisterSize(uint32_t type);
+
     size_t calculateThreadSlotOffset(EuThread::ThreadId threadId);
     size_t calculateRegisterOffsetInThreadSlot(const SIP::regset_desc *const regdesc, uint32_t start);
 

level_zero/tools/test/unit_tests/sources/debug/debug_session_tests.cpp

@@ -36,6 +36,7 @@ struct MockDebugSession : public L0::DebugSessionImp {
     using L0::DebugSessionImp::generateEventsAndResumeStoppedThreads;
     using L0::DebugSessionImp::generateEventsForPendingInterrupts;
     using L0::DebugSessionImp::generateEventsForStoppedThreads;
+    using L0::DebugSessionImp::getRegisterSize;
     using L0::DebugSessionImp::getStateSaveAreaHeader;
     using L0::DebugSessionImp::markPendingInterruptsOrAddToNewlyStoppedFromRaisedAttention;
     using L0::DebugSessionImp::newAttentionRaised;
@@ -1717,6 +1718,17 @@ TEST_F(DebugSessionRegistersAccessTest, givenTypeToRegsetDescCalledThenCorrectRe
     EXPECT_NE(session->typeToRegsetDesc(ZET_DEBUG_REGSET_TYPE_SBA_INTEL_GPU), nullptr);
 }
 
+TEST_F(DebugSessionRegistersAccessTest, givenValidRegisterWhenGettingSizeThenCorrectSizeIsReturned) {
+    session->stateSaveAreaHeader = MockSipData::createStateSaveAreaHeader(2);
+    auto pStateSaveAreaHeader = session->getStateSaveAreaHeader();
+    EXPECT_EQ(pStateSaveAreaHeader->regHeader.grf.bytes, session->getRegisterSize(ZET_DEBUG_REGSET_TYPE_GRF_INTEL_GPU));
+}
+
+TEST_F(DebugSessionRegistersAccessTest, givenInvalidRegisterWhenGettingSizeThenZeroSizeIsReturned) {
+    session->stateSaveAreaHeader = MockSipData::createStateSaveAreaHeader(2);
+    EXPECT_EQ(0u, session->getRegisterSize(ZET_DEBUG_REGSET_TYPE_INVALID_INTEL_GPU));
+}
+
 TEST_F(DebugSessionRegistersAccessTest, givenUnsupportedRegisterTypeWhenReadRegistersCalledThenErrorInvalidArgumentIsReturned) {
     session->areRequestedThreadsStoppedReturnValue = 1;
     session->stateSaveAreaHeader = MockSipData::createStateSaveAreaHeader(2);