Initial work on power assertions

CMakeLists.txt

@@ -39,6 +39,8 @@ set(iokitd_sources
 	src/IORegistryEntry.mm
 	src/IODisplayConnect.mm
 	src/IODisplayConnectX11.mm
+	src/PowerAssertions.mm
+	src/PowerAssertionsX11.mm
 	${CMAKE_CURRENT_BINARY_DIR}/iokitmigServer.c
 	${CMAKE_CURRENT_BINARY_DIR}/powermanagementServer.c
 	${CMAKE_CURRENT_SOURCE_DIR}/../IOKitUser/IOCFSerialize.c

src/PowerAssertions.h

@@ -0,0 +1,46 @@
+#ifndef POWERASSERTIONS_H
+#define POWERASSERTIONS_H
+#include <IOKit/pwr_mgt/IOPMLib.h>
+#include <unordered_map>
+#include <dispatch/dispatch.h>
+
+void initPM();
+
+class PowerAssertion
+{
+public:
+	// How many AppAssertions instances reference this assertion (*not* a sum of references within all these AppAssertions instances)
+	void addRef();
+	void delRef();
+protected:
+	virtual void activate() = 0;
+	virtual void deactivate() = 0;
+private:
+	int m_numHoldingApps = 0;
+};
+
+class AppAssertions
+{
+private:
+	AppAssertions(int pid);
+	~AppAssertions();
+public:
+	static AppAssertions* get(int pid);
+
+	struct HeldAssertion
+	{
+		int refcount;
+		PowerAssertion* assertion;
+	};
+	HeldAssertion* getAssertion(int assertionId);
+	int createAssertion(PowerAssertion* which);
+	void killAssertion(int assertionId);
+private:
+	std::unordered_map<int, HeldAssertion> m_heldAssertions;
+	int m_nextAssertionNumber = 1;
+
+	dispatch_source_t m_processSource;
+	static std::unordered_map<int, AppAssertions*> m_instances;
+};
+
+#endif

src/PowerAssertions.mm

@@ -0,0 +1,214 @@
+#include "PowerAssertions.h"
+#include "PowerAssertionsX11.h"
+#include <stdexcept>
+#include <iostream>
+#include <string>
+#include <CoreFoundation/CoreFoundation.h>
+#include <bsm/libbsm.h>
+#import <Foundation/NSString.h>
+
+extern "C" {
+#include "powermanagementServer.h"
+}
+
+std::unordered_map<int, AppAssertions*> AppAssertions::m_instances;
+static std::unordered_map<std::string, PowerAssertion*> g_assertionKinds;
+
+void initPM()
+{
+	NSString* str = (NSString*) kIOPMAssertPreventUserIdleDisplaySleep;
+	g_assertionKinds.insert(std::make_pair([str UTF8String], new PowerAssertionPreventDisplaySleep));
+}
+
+static PowerAssertion* getByType(NSString* type)
+{
+	auto it = g_assertionKinds.find([type UTF8String]);
+	if (it != g_assertionKinds.end())
+		return it->second;
+	return nullptr;
+}
+
+void PowerAssertion::addRef()
+{
+	if (m_numHoldingApps == 0)
+		activate();
+	m_numHoldingApps++;
+}
+
+void PowerAssertion::delRef()
+{
+	m_numHoldingApps--;
+	if (m_numHoldingApps == 0)
+		deactivate();
+}
+
+AppAssertions::AppAssertions(int pid)
+{
+	m_processSource = dispatch_source_create(DISPATCH_SOURCE_TYPE_PROC, pid, DISPATCH_PROC_EXIT, dispatch_get_main_queue());
+	if (!m_processSource)
+		throw std::runtime_error("dispatch_source_create(DISPATCH_SOURCE_TYPE_PROC) failed");
+
+	if (m_processSource)
+	{
+		dispatch_source_set_event_handler(m_processSource, ^{
+			m_instances.erase(pid);
+			delete this;
+		});
+		dispatch_resume(m_processSource);
+	}
+}
+
+AppAssertions::~AppAssertions()
+{
+	for (auto& [assertionId, ha] : m_heldAssertions)
+		ha.assertion->delRef();
+		
+	if (m_processSource != nullptr)
+		dispatch_release(m_processSource);
+}
+
+AppAssertions* AppAssertions::get(int pid)
+{
+	auto it = m_instances.find(pid);
+	if (it != m_instances.end())
+		return it->second;
+	
+	try
+	{
+		AppAssertions* aa = new AppAssertions(pid);
+		m_instances.insert(std::make_pair(pid, aa));
+		return aa;
+	}
+	catch (const std::exception& e)
+	{
+		std::cerr << "AppAssertions::get() failed for PID " << pid << ": " << e.what() << std::endl;
+		return nullptr;
+	}
+}
+
+AppAssertions::HeldAssertion* AppAssertions::getAssertion(int assertionId)
+{
+	auto it = m_heldAssertions.find(assertionId);
+	if (it != m_heldAssertions.end())
+		return &it->second;
+	return nullptr;
+}
+
+int AppAssertions::createAssertion(PowerAssertion* which)
+{
+	int number = m_nextAssertionNumber++;
+	HeldAssertion held = { 1, which };
+	m_heldAssertions.insert(std::make_pair(number, held));
+	which->addRef();
+	return number;
+}
+
+void AppAssertions::killAssertion(int assertionId)
+{
+	auto it = m_heldAssertions.find(assertionId);
+	if (it == m_heldAssertions.end())
+		return;
+
+	it->second.assertion->delRef();
+	m_heldAssertions.erase(it);
+}
+
+kern_return_t _io_pm_assertion_create
+(
+	mach_port_t server,
+	audit_token_t token,
+	vm_offset_t props,
+	mach_msg_type_number_t propsCnt,
+	int *assertion_id,
+	int *disableAppSleep,
+	int *enTrIntensity,
+	int *return_code
+)
+{
+	if (!props)
+		return KERN_INVALID_ARGUMENT;
+
+	CFDataRef data = CFDataCreateWithBytesNoCopy(nullptr, (const UInt8*) props, propsCnt, kCFAllocatorNull);
+	CFDictionaryRef properties = (CFDictionaryRef) CFPropertyListCreateWithData(nullptr, data, kCFPropertyListImmutable, nullptr, nullptr);
+
+	CFRelease(data);
+
+	vm_deallocate(mach_task_self(), props, propsCnt);
+
+	if (!properties)
+		return KERN_INVALID_ARGUMENT;
+	
+	if (CFGetTypeID(properties) != CFDictionaryGetTypeID())
+	{
+		std::cerr << "Unexpected properties data type\n";
+		CFRelease(properties);
+		return KERN_INVALID_ARGUMENT;
+	}
+
+	CFStringRef cftype = (CFStringRef) CFDictionaryGetValue(properties, kIOPMAssertionTypeKey);
+
+	if (!cftype || CFStringGetTypeID() != CFGetTypeID(cftype))
+	{
+		std::cerr << "kIOPMAssertionTypeKey is not a string\n";
+		CFRelease(properties);
+		return KERN_INVALID_ARGUMENT;
+	}
+
+	PowerAssertion* pa = getByType((NSString*) cftype);
+	if (pa == nullptr)
+	{
+		std::cerr << "Unsupported assertion type\n";
+		CFRelease(properties);
+		return KERN_INVALID_ARGUMENT;
+	}
+
+	pid_t callerPid;
+	audit_token_to_au32(token, nullptr, nullptr, nullptr, nullptr, nullptr, &callerPid, nullptr, nullptr);
+
+	std::cout << "Creating assertion for PID " << callerPid << std::endl;
+	*assertion_id = AppAssertions::get(callerPid)->createAssertion(pa);
+	*return_code = 0;
+	
+    return KERN_SUCCESS;
+}
+
+kern_return_t _io_pm_assertion_retain_release
+(
+	mach_port_t server,
+	audit_token_t token,
+	int assertion_id,
+	int action,
+	int *retainCnt,
+	int *disableAppSleep,
+	int *enableAppSleep,
+	int *return_code
+)
+{
+	pid_t callerPid;
+	audit_token_to_au32(token, nullptr, nullptr, nullptr, nullptr, nullptr, &callerPid, nullptr, nullptr);
+
+	AppAssertions* appAssertions = AppAssertions::get(callerPid);
+	AppAssertions::HeldAssertion* ha = appAssertions->getAssertion(assertion_id);
+
+	if (!ha)
+	{
+		std::cerr << "_io_pm_assertion_retain_release() cannot find assertion " << assertion_id << std::endl;
+		return KERN_INVALID_ARGUMENT;
+	}
+
+	*return_code = 0;
+	if (action == kIOPMAssertionMIGDoRetain)
+	{
+		*retainCnt = ++ha->refcount;
+	}
+	else if (action == kIOPMAssertionMIGDoRelease)
+	{
+		*retainCnt = --ha->refcount;
+		if (ha->refcount == 0)
+			appAssertions->killAssertion(assertion_id);
+	}
+	else
+		return KERN_INVALID_ARGUMENT;
+
+    return KERN_NOT_SUPPORTED;
+}

src/PowerAssertionsX11.h

@@ -0,0 +1,13 @@
+#ifndef POWERASSERTIONSX11_H
+#define POWERASSERTIONSX11_H
+#include "PowerAssertions.h"
+
+class PowerAssertionPreventDisplaySleep : public PowerAssertion
+{
+protected:
+	void activate() override;
+	void deactivate() override;
+};
+
+#endif
+

src/PowerAssertionsX11.mm

@@ -0,0 +1,12 @@
+#include "PowerAssertionsX11.h"
+#include <iostream>
+
+void PowerAssertionPreventDisplaySleep::activate()
+{
+	std::cout << "STUB: PowerAssertionPreventDisplaySleep::activate()\n";
+}
+
+void PowerAssertionPreventDisplaySleep::deactivate()
+{
+	std::cout << "STUB: PowerAssertionPreventDisplaySleep::deactivate()\n";
+}

src/main.mm

@@ -4,11 +4,13 @@
 #include <dispatch/dispatch.h>
 #include <dispatch/private.h>
 #include <IOKit/IOKitKeys.h>
+#include <IOKit/pwr_mgt/IOPMLibPrivate.h>
 #include <cstdlib>
 #include "iokitd.h"
 #include "iokitmig.h"
 #include "IOObject.h"
 #include "IODisplayConnectX11.h"
+#include "PowerAssertions.h"
 
 extern "C" {
 #include "iokitmigServer.h"
@@ -33,7 +35,7 @@ int main(int argc, const char** argv)
 		return 1;
 	}
 
-	ret = bootstrap_check_in(bootstrap_port, SERVICE_NAME, &g_powerManagementPort);
+	ret = bootstrap_check_in(bootstrap_port, kIOPMServerBootstrapName, &g_powerManagementPort);
 
 	if (ret != KERN_SUCCESS)
 	{
@@ -50,6 +52,7 @@ int main(int argc, const char** argv)
 
 	// Build IOKit registry here
 	discoverAllDevices();
+	initPM();
 
 	dispatch_queue_t queue = dispatch_get_main_queue();
 
@@ -100,6 +103,7 @@ int main(int argc, const char** argv)
 
 	dispatch_resume(portSource);
 	dispatch_resume(deathSource);
+	dispatch_resume(pwrMgmtPortSource);
 
 	os_log(OS_LOG_DEFAULT, "iokitd up and running.");
 

src/stubs-powermanagement.c

@@ -131,25 +131,6 @@ kern_return_t _io_pm_last_wake_time
 }
 
 
-/* Routine io_pm_assertion_create */
-
-kern_return_t _io_pm_assertion_create
-(
-	mach_port_t server,
-	audit_token_t token,
-	vm_offset_t props,
-	mach_msg_type_number_t propsCnt,
-	int *assertion_id,
-	int *disableAppSleep,
-	int *enTrIntensity,
-	int *return_code
-)
-{
-    STUB();
-    return KERN_NOT_SUPPORTED;
-}
-
-
 /* Routine io_pm_assertion_set_properties */
 
 kern_return_t _io_pm_assertion_set_properties
@@ -168,26 +149,6 @@ kern_return_t _io_pm_assertion_set_properties
     return KERN_NOT_SUPPORTED;
 }
 
-
-/* Routine io_pm_assertion_retain_release */
-
-kern_return_t _io_pm_assertion_retain_release
-(
-	mach_port_t server,
-	audit_token_t token,
-	int assertion_id,
-	int action,
-	int *retainCnt,
-	int *disableAppSleep,
-	int *enableAppSleep,
-	int *return_code
-)
-{
-    STUB();
-    return KERN_NOT_SUPPORTED;
-}
-
-
 /* Routine io_pm_assertion_copy_details */
 
 kern_return_t _io_pm_assertion_copy_details