From cf555cdf56f0cf619d269ac08bd16025eba571dc Mon Sep 17 00:00:00 2001
From: HaoYan Qi <haoyanqi666@gmail.com>
Date: Thu, 16 Nov 2023 16:01:08 +1100
Subject: [PATCH] ensure PcmMsrClient can only opened by root

---
 src/MacMSRDriver/PcmMsr/PcmMsrClient.cpp | 15 ++++++++++++++-
 src/MacMSRDriver/PcmMsr/PcmMsrClient.h   |  4 +++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/MacMSRDriver/PcmMsr/PcmMsrClient.cpp b/src/MacMSRDriver/PcmMsr/PcmMsrClient.cpp
index cfca0b0e..92f0bbd3 100644
--- a/src/MacMSRDriver/PcmMsr/PcmMsrClient.cpp
+++ b/src/MacMSRDriver/PcmMsr/PcmMsrClient.cpp
@@ -41,10 +41,23 @@ IOReturn PcmMsrClientClassName::externalMethod(uint32_t selector, IOExternalMeth
 	return super::externalMethod(selector, args, dispatch, target, reference);
 }
 
+bool PcmMsrClientClassName::initWithTask(task_t owningTask, void *securityToken, UInt32 type, OSDictionary *properties)
+{
+    if(!IOUserClient::initWithTask(owningTask, securityToken, type, properties)) {
+        return false;
+    }
+    
+    sSecurityToken = securityToken;
+    return true;
+}
+
 bool PcmMsrClientClassName::start(IOService* provider)
 {
 	bool result = false;
-    
+
+    if(clientHasPrivilege(sSecurityToken, kIOClientPrivilegeAdministrator) != kIOReturnSuccess) 
+		return false;
+		
     fProvider = OSDynamicCast(PcmMsrDriverClassName, provider);
     
     if (fProvider != NULL) {
diff --git a/src/MacMSRDriver/PcmMsr/PcmMsrClient.h b/src/MacMSRDriver/PcmMsr/PcmMsrClient.h
index 54281e54..77da0b93 100644
--- a/src/MacMSRDriver/PcmMsr/PcmMsrClient.h
+++ b/src/MacMSRDriver/PcmMsr/PcmMsrClient.h
@@ -14,11 +14,13 @@ class PcmMsrClientClassName : public IOUserClient
 
 protected:
     PcmMsrDriverClassName*                  fProvider;
+    void*                                   sSecurityToken;
     static const IOExternalMethodDispatch   sMethods[kNumberOfMethods];
 
 public:
+    virtual bool initWithTask(task_t owningTask, void *securityToken, UInt32 type, OSDictionary *properties) override;
     virtual bool start(IOService *provider) override;
-
+    
     virtual IOReturn clientClose(void) override;
 
     virtual bool didTerminate(IOService* provider, IOOptionBits opts, bool* defer) override;