From 569d580615604c74c4420d0d36012518adf81c92 Mon Sep 17 00:00:00 2001
From: dubj <dubingjian@huawei.com>
Date: Thu, 26 May 2022 11:19:25 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=88=AA=E5=B1=8F=E6=9D=83?=
 =?UTF-8?q?=E9=99=90=E6=9C=8D=E5=8A=A1=E7=AB=AF=E9=AA=8C=E8=AF=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: dubj <dubingjian@huawei.com>
Change-Id: I11ee0a577f369a7b2d7121161e220ae6f3c04678
---
 dmserver/BUILD.gn                             |  2 ++
 dmserver/include/display_manager_stub.h       |  3 ++
 dmserver/src/display_manager_stub.cpp         | 36 ++++++++++++++++---
 .../screenshot/native_screenshot_module.cpp   |  7 ----
 4 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/dmserver/BUILD.gn b/dmserver/BUILD.gn
index ca7fdfa212..57c0c61f36 100644
--- a/dmserver/BUILD.gn
+++ b/dmserver/BUILD.gn
@@ -51,6 +51,8 @@ ohos_shared_library("libdms") {
 
   external_deps = [
     "bytrace_standard:bytrace_core",
+    "ability_runtime:ability_manager",
+    "access_token:libaccesstoken_sdk",
     "graphic_standard:surface",
     "hilog_native:libhilog",
     "ipc:ipc_core",
diff --git a/dmserver/include/display_manager_stub.h b/dmserver/include/display_manager_stub.h
index d88c4a0c65..32488a7c2a 100644
--- a/dmserver/include/display_manager_stub.h
+++ b/dmserver/include/display_manager_stub.h
@@ -27,6 +27,9 @@ public:
     ~DisplayManagerStub() = default;
     virtual int32_t OnRemoteRequest(uint32_t code, MessageParcel &data, MessageParcel &reply,
         MessageOption &option) override;
+private:
+    bool CheckCallingPermission(const std::string &permission);
+    bool IsStartByHdcd();
 };
 } // namespace OHOS::Rosen
 
diff --git a/dmserver/src/display_manager_stub.cpp b/dmserver/src/display_manager_stub.cpp
index e7a520e8fd..b4e0537aca 100644
--- a/dmserver/src/display_manager_stub.cpp
+++ b/dmserver/src/display_manager_stub.cpp
@@ -19,6 +19,8 @@
 
 #include <ipc_skeleton.h>
 
+#include "accesstoken_kit.h"
+#include "bundle_constants.h"
 #include "window_manager_hilog.h"
 
 #include "transaction/rs_interfaces.h"
@@ -28,6 +30,27 @@ namespace {
     constexpr HiviewDFX::HiLogLabel LABEL = {LOG_CORE, HILOG_DOMAIN_DISPLAY, "DisplayManagerStub"};
 }
 
+bool DisplayManagerStub::CheckCallingPermission(const std::string &permission)
+{
+    if (Security::AccessToken::AccessTokenKit::VerifyAccessToken(IPCSkeleton::GetCallingTokenID(), permission) !=
+        AppExecFwk::Constants::PERMISSION_GRANTED) {
+        return false;
+    }
+    return true;
+}
+
+bool DisplayManagerStub::IsStartByHdcd()
+{
+    OHOS::Security::AccessToken::NativeTokenInfo info;
+    if (Security::AccessToken::AccessTokenKit::GetNativeTokenInfo(IPCSkeleton::GetCallingTokenID(), info) != 0) {
+        return false;
+    }
+    if (info.processName.compare("hdcd") == 0) {
+        return true;
+    }
+    return false;
+}
+
 int32_t DisplayManagerStub::OnRemoteRequest(uint32_t code, MessageParcel &data, MessageParcel &reply,
     MessageOption &option)
 {
@@ -107,9 +130,14 @@ int32_t DisplayManagerStub::OnRemoteRequest(uint32_t code, MessageParcel &data,
             break;
         }
         case TRANS_ID_GET_DISPLAY_SNAPSHOT: {
-            DisplayId displayId = data.ReadUint64();
-            std::shared_ptr<Media::PixelMap> displaySnapshot = GetDisplaySnapshot(displayId);
-            reply.WriteParcelable(displaySnapshot == nullptr ? nullptr : displaySnapshot.get());
+            if (CheckCallingPermission("ohos.permission.CAPTURE_SCREEN") || IsStartByHdcd()) {
+                DisplayId displayId = data.ReadUint64();
+                std::shared_ptr<Media::PixelMap> displaySnapshot = GetDisplaySnapshot(displayId);
+                reply.WriteParcelable(displaySnapshot == nullptr ? nullptr : displaySnapshot.get());
+                break;
+            }
+            WLOGFE("check permission failed!");
+            reply.WriteParcelable(nullptr);
             break;
         }
         case TRANS_ID_REGISTER_DISPLAY_MANAGER_AGENT: {
@@ -326,4 +354,4 @@ int32_t DisplayManagerStub::OnRemoteRequest(uint32_t code, MessageParcel &data,
     }
     return 0;
 }
-} // namespace OHOS::Rosen
\ No newline at end of file
+} // namespace OHOS::Rosen
diff --git a/interfaces/kits/napi/screenshot/native_screenshot_module.cpp b/interfaces/kits/napi/screenshot/native_screenshot_module.cpp
index b800fe155e..fbb32433af 100644
--- a/interfaces/kits/napi/screenshot/native_screenshot_module.cpp
+++ b/interfaces/kits/napi/screenshot/native_screenshot_module.cpp
@@ -164,13 +164,6 @@ static void GetScreenshotParam(napi_env env, std::unique_ptr<Param> &param, napi
 
 static void AsyncGetScreenshot(napi_env env, std::unique_ptr<Param> &param)
 {
-    if (!CheckCallingPermission("ohos.permission.CAPTURE_SCREEN")) {
-        WLOGFE("Get Screenshot failed. Do not have permission!");
-        param->image = nullptr;
-        param->wret = WMError::WM_ERROR_INVALID_PERMISSION;
-        param->errMessage = "Get Screenshot Failed: Do not have ohos.permission.CAPTURE_SCREEN";
-        return;
-    }
     if (!param->validInputParam) {
         WLOGFE("Invalid Input Param!");
         param->image = nullptr;
-- 
Gitee