Merge pull request #1843 from AllanZyne/review/yang/invalid_arguments

omarahmed1111 · web-flow · commit d8058ed3f191 · 2024-08-07T12:30:47.000+01:00
[DeviceSanitizer] Support check invalid kernel argument
diff --git a/source/loader/CMakeLists.txt b/source/loader/CMakeLists.txt
@@ -142,6 +142,8 @@ if(UR_ENABLE_SANITIZER)
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_report.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_shadow_setup.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_shadow_setup.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_validator.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_validator.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/common.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/stacktrace.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/stacktrace.hpp
diff --git a/source/loader/layers/sanitizer/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan_interceptor.cpp
@@ -16,6 +16,7 @@
 #include "asan_quarantine.hpp"
 #include "asan_report.hpp"
 #include "asan_shadow_setup.hpp"
+#include "asan_validator.hpp"
 #include "stacktrace.hpp"
 #include "ur_sanitizer_utils.hpp"
 
@@ -615,6 +616,9 @@ SanitizerInterceptor::insertDevice(ur_device_handle_t Device,
 
     DI = std::make_shared<ur_sanitizer_layer::DeviceInfo>(Device);
 
+    DI->IsSupportSharedSystemUSM = GetDeviceUSMCapability(
+        Device, UR_DEVICE_INFO_USM_SYSTEM_SHARED_SUPPORT);
+
     // Query alignment
     UR_CALL(getContext()->urDdiTable.Device.pfnGetInfo(
         Device, UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN, sizeof(DI->Alignment),
@@ -683,8 +687,25 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
     auto Program = GetProgram(Kernel);
 
     do {
-        // Set membuffer arguments
         auto KernelInfo = getKernelInfo(Kernel);
+
+        // Validate pointer arguments
+        if (Options(logger).DetectKernelArguments) {
+            for (const auto &[ArgIndex, PtrPair] : KernelInfo->PointerArgs) {
+                auto Ptr = PtrPair.first;
+                if (Ptr == nullptr) {
+                    continue;
+                }
+                if (auto ValidateResult = ValidateUSMPointer(
+                        Context, DeviceInfo->Handle, (uptr)Ptr)) {
+                    ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr,
+                                                ValidateResult, PtrPair.second);
+                    exit(1);
+                }
+            }
+        }
+
+        // Set membuffer arguments
         for (const auto &[ArgIndex, MemBuffer] : KernelInfo->BufferArgs) {
             char *ArgPointer = nullptr;
             UR_CALL(MemBuffer->getHandle(DeviceInfo->Handle, ArgPointer));
diff --git a/source/loader/layers/sanitizer/asan_interceptor.hpp b/source/loader/layers/sanitizer/asan_interceptor.hpp
@@ -41,6 +41,9 @@ struct DeviceInfo {
     uptr ShadowOffset = 0;
     uptr ShadowOffsetEnd = 0;
 
+    // Device features
+    bool IsSupportSharedSystemUSM;
+
     ur_mutex Mutex;
     std::queue<std::shared_ptr<AllocInfo>> Quarantine;
     size_t QuarantineSize = 0;
@@ -78,6 +81,8 @@ struct KernelInfo {
     ur_shared_mutex Mutex;
     std::atomic<int32_t> RefCount = 1;
     std::unordered_map<uint32_t, std::shared_ptr<MemBuffer>> BufferArgs;
+    std::unordered_map<uint32_t, std::pair<const void *, StackTrace>>
+        PointerArgs;
 
     // Need preserve the order of local arguments
     std::map<uint32_t, LocalArgsInfo> LocalArgs;
diff --git a/source/loader/layers/sanitizer/asan_options.hpp b/source/loader/layers/sanitizer/asan_options.hpp
@@ -38,6 +38,7 @@ struct AsanOptions {
     uint32_t MaxQuarantineSizeMB = 0;
     bool DetectLocals = true;
     bool DetectPrivates = true;
+    bool DetectKernelArguments = true;
 
   private:
     AsanOptions(logger::Logger &logger) {
@@ -93,6 +94,7 @@ struct AsanOptions {
         SetBoolOption("debug", Debug);
         SetBoolOption("detect_locals", DetectLocals);
         SetBoolOption("detect_privates", DetectPrivates);
+        SetBoolOption("detect_kernel_arguments", DetectKernelArguments);
 
         auto KV = OptionsEnvMap->find("quarantine_size_mb");
         if (KV != OptionsEnvMap->end()) {
diff --git a/source/loader/layers/sanitizer/asan_report.cpp b/source/loader/layers/sanitizer/asan_report.cpp
@@ -11,16 +11,32 @@
  */
 
 #include "asan_report.hpp"
-#include "asan_options.hpp"
-
 #include "asan_allocator.hpp"
 #include "asan_interceptor.hpp"
 #include "asan_libdevice.hpp"
+#include "asan_options.hpp"
+#include "asan_validator.hpp"
 #include "ur_sanitizer_layer.hpp"
 #include "ur_sanitizer_utils.hpp"
 
 namespace ur_sanitizer_layer {
 
+namespace {
+
+void PrintAllocateInfo(uptr Addr, const AllocInfo *AI) {
+    getContext()->logger.always("{} is located inside of {} region [{}, {})",
+                                (void *)Addr, ToString(AI->Type),
+                                (void *)AI->UserBegin, (void *)AI->UserEnd);
+    getContext()->logger.always("allocated here:");
+    AI->AllocStack.print();
+    if (AI->IsReleased) {
+        getContext()->logger.always("freed here:");
+        AI->ReleaseStack.print();
+    }
+}
+
+} // namespace
+
 void ReportBadFree(uptr Addr, const StackTrace &stack,
                    const std::shared_ptr<AllocInfo> &AI) {
     getContext()->logger.always(
@@ -34,11 +50,7 @@ void ReportBadFree(uptr Addr, const StackTrace &stack,
 
     assert(AI && !AI->IsReleased && "Chunk must be not released");
 
-    getContext()->logger.always("{} is located inside of {} region [{}, {})",
-                                (void *)Addr, ToString(AI->Type),
-                                (void *)AI->UserBegin, (void *)AI->UserEnd);
-    getContext()->logger.always("allocated here:");
-    AI->AllocStack.print();
+    PrintAllocateInfo(Addr, AI.get());
 }
 
 void ReportBadContext(uptr Addr, const StackTrace &stack,
@@ -48,16 +60,7 @@ void ReportBadContext(uptr Addr, const StackTrace &stack,
         (void *)Addr);
     stack.print();
 
-    getContext()->logger.always("{} is located inside of {} region [{}, {})",
-                                (void *)Addr, ToString(AI->Type),
-                                (void *)AI->UserBegin, (void *)AI->UserEnd);
-    getContext()->logger.always("allocated here:");
-    AI->AllocStack.print();
-
-    if (AI->IsReleased) {
-        getContext()->logger.always("freed here:");
-        AI->ReleaseStack.print();
-    }
+    PrintAllocateInfo(Addr, AI.get());
 }
 
 void ReportDoubleFree(uptr Addr, const StackTrace &Stack,
@@ -139,16 +142,10 @@ void ReportUseAfterFree(const DeviceSanitizerReport &Report,
                     "Failed to find which chunck {} is allocated",
                     (void *)Report.Address);
             }
-            assert(AllocInfo->IsReleased);
+            assert(AllocInfo->IsReleased &&
+                   "It must be released since it's use-after-free");
 
-            getContext()->logger.always(
-                "{} is located inside of {} region [{}, {})",
-                (void *)Report.Address, ToString(AllocInfo->Type),
-                (void *)AllocInfo->UserBegin, (void *)AllocInfo->UserEnd);
-            getContext()->logger.always("allocated here:");
-            AllocInfo->AllocStack.print();
-            getContext()->logger.always("released here:");
-            AllocInfo->ReleaseStack.print();
+            PrintAllocateInfo(Report.Address, AllocInfo.get());
         }
     } else {
         getContext()->logger.always(
@@ -157,4 +154,47 @@ void ReportUseAfterFree(const DeviceSanitizerReport &Report,
     }
 }
 
+void ReportInvalidKernelArgument(ur_kernel_handle_t Kernel, uint32_t ArgIndex,
+                                 uptr Addr, const ValidateUSMResult &VR,
+                                 StackTrace Stack) {
+    getContext()->logger.always("\n====ERROR: DeviceSanitizer: "
+                                "invalid-argument on kernel <{}>",
+                                DemangleName(GetKernelName(Kernel)));
+    Stack.print();
+    auto &AI = VR.AI;
+    switch (VR.Type) {
+    case ValidateUSMResult::MAYBE_HOST_POINTER:
+        getContext()->logger.always("The {}th argument {} is not a USM pointer",
+                                    ArgIndex + 1, (void *)Addr);
+        break;
+    case ValidateUSMResult::RELEASED_POINTER:
+        getContext()->logger.always(
+            "The {}th argument {} is a released USM pointer", ArgIndex,
+            (void *)Addr);
+        PrintAllocateInfo(Addr, AI.get());
+        break;
+    case ValidateUSMResult::BAD_CONTEXT:
+        getContext()->logger.always(
+            "The {}th argument {} is allocated in other context", ArgIndex,
+            (void *)Addr);
+        PrintAllocateInfo(Addr, AI.get());
+        break;
+    case ValidateUSMResult::BAD_DEVICE:
+        getContext()->logger.always(
+            "The {}th argument {} is allocated in other device", ArgIndex,
+            (void *)Addr);
+        PrintAllocateInfo(Addr, AI.get());
+        break;
+    case ValidateUSMResult::OUT_OF_BOUNDS:
+        getContext()->logger.always(
+            "The {}th argument {} is located outside of its region [{}, {})",
+            ArgIndex, (void *)Addr, (void *)AI->UserBegin, (void *)AI->UserEnd);
+        getContext()->logger.always("allocated here:");
+        AI->AllocStack.print();
+        break;
+    default:
+        break;
+    }
+}
+
 } // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/asan_report.hpp b/source/loader/layers/sanitizer/asan_report.hpp
@@ -21,6 +21,7 @@ namespace ur_sanitizer_layer {
 struct DeviceSanitizerReport;
 struct AllocInfo;
 struct StackTrace;
+struct ValidateUSMResult;
 
 void ReportBadFree(uptr Addr, const StackTrace &stack,
                    const std::shared_ptr<AllocInfo> &AllocInfo);
@@ -40,4 +41,8 @@ void ReportGenericError(const DeviceSanitizerReport &Report,
 void ReportUseAfterFree(const DeviceSanitizerReport &Report,
                         ur_kernel_handle_t Kernel, ur_context_handle_t Context);
 
+void ReportInvalidKernelArgument(ur_kernel_handle_t Kernel, uint32_t ArgIndex,
+                                 uptr Addr, const ValidateUSMResult &VR,
+                                 StackTrace Stack);
+
 } // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/asan_validator.cpp b/source/loader/layers/sanitizer/asan_validator.cpp
@@ -0,0 +1,77 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ * @file asan_validator.cpp
+ *
+ */
+
+#include "asan_validator.hpp"
+#include "asan_interceptor.hpp"
+#include "ur_sanitizer_utils.hpp"
+
+namespace ur_sanitizer_layer {
+
+namespace {
+
+bool IsSameDevice(ur_device_handle_t Device1, ur_device_handle_t Device2) {
+    if (Device1 == Device2) {
+        return true;
+    }
+    auto RootDevice1 = GetParentDevice(Device1);
+    RootDevice1 = RootDevice1 ? RootDevice1 : Device1;
+    auto RootDevice2 = GetParentDevice(Device2);
+    RootDevice2 = RootDevice2 ? RootDevice2 : Device2;
+    if (RootDevice1 == RootDevice2) {
+        return true;
+    }
+    return false;
+}
+
+} // namespace
+
+ValidateUSMResult ValidateUSMPointer(ur_context_handle_t Context,
+                                     ur_device_handle_t Device, uptr Ptr) {
+    assert(Ptr != 0 && "Don't validate nullptr here");
+
+    auto AllocInfoItOp = getContext()->interceptor->findAllocInfoByAddress(Ptr);
+    if (!AllocInfoItOp) {
+        auto DI = getContext()->interceptor->getDeviceInfo(Device);
+        bool IsSupportSharedSystemUSM = DI->IsSupportSharedSystemUSM;
+        if (IsSupportSharedSystemUSM) {
+            // maybe it's host pointer
+            return ValidateUSMResult::success();
+        }
+        return ValidateUSMResult::fail(ValidateUSMResult::MAYBE_HOST_POINTER);
+    }
+
+    auto AllocInfo = AllocInfoItOp.value()->second;
+
+    if (AllocInfo->Context != Context) {
+        return ValidateUSMResult::fail(ValidateUSMResult::BAD_CONTEXT,
+                                       AllocInfo);
+    }
+
+    if (AllocInfo->Device && !IsSameDevice(AllocInfo->Device, Device)) {
+        return ValidateUSMResult::fail(ValidateUSMResult::BAD_DEVICE,
+                                       AllocInfo);
+    }
+
+    if (AllocInfo->IsReleased) {
+        return ValidateUSMResult::fail(ValidateUSMResult::RELEASED_POINTER,
+                                       AllocInfo);
+    }
+
+    if (Ptr < AllocInfo->UserBegin || Ptr >= AllocInfo->UserEnd) {
+        return ValidateUSMResult::fail(ValidateUSMResult::OUT_OF_BOUNDS,
+                                       AllocInfo);
+    }
+
+    return ValidateUSMResult::success();
+}
+
+} // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/asan_validator.hpp b/source/loader/layers/sanitizer/asan_validator.hpp
@@ -0,0 +1,50 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ * @file asan_validator.hpp
+ *
+ */
+#pragma once
+
+#include "asan_allocator.hpp"
+
+namespace ur_sanitizer_layer {
+
+struct ValidateUSMResult {
+    enum ErrorType {
+        SUCCESS,
+        NULL_POINTER,
+        MAYBE_HOST_POINTER,
+        RELEASED_POINTER,
+        BAD_CONTEXT,
+        BAD_DEVICE,
+        OUT_OF_BOUNDS
+    };
+    ErrorType Type;
+    std::shared_ptr<AllocInfo> AI;
+
+    operator bool() { return Type != SUCCESS; }
+
+    static ValidateUSMResult success() { return {SUCCESS, nullptr}; }
+
+    static ValidateUSMResult fail(ErrorType Type,
+                                  const std::shared_ptr<AllocInfo> &AI) {
+        assert(Type != SUCCESS && "The error type shouldn't be SUCCESS");
+        return {Type, AI};
+    }
+
+    static ValidateUSMResult fail(ErrorType Type) {
+        assert(Type != SUCCESS && "The error type shouldn't be SUCCESS");
+        return {Type, nullptr};
+    }
+};
+
+ValidateUSMResult ValidateUSMPointer(ur_context_handle_t Context,
+                                     ur_device_handle_t Device, uptr Ptr);
+
+} // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/ur_sanddi.cpp b/source/loader/layers/sanitizer/ur_sanddi.cpp
diff --git a/source/loader/layers/sanitizer/ur_sanitizer_utils.cpp b/source/loader/layers/sanitizer/ur_sanitizer_utils.cpp
diff --git a/source/loader/layers/sanitizer/ur_sanitizer_utils.hpp b/source/loader/layers/sanitizer/ur_sanitizer_utils.hpp
