refactor: extract some desktop-related APIs to an independent class

Signed-off-by: Trumeet <17158086+Trumeet@users.noreply.github.com>

3 files changed, 115 insertions, 33 deletions

diff --git a/src/mingwMain/kotlin/moe.yuuta.desktop/Desktop.kt b/src/mingwMain/kotlin/moe.yuuta.desktop/Desktop.kt
index 9d881c9..1d64fb8 100644
--- a/src/mingwMain/kotlin/moe.yuuta.desktop/Desktop.kt
+++ b/src/mingwMain/kotlin/moe.yuuta.desktop/Desktop.kt
@@ -9,20 +9,24 @@ import platform.windows.*
 import kotlin.native.concurrent.TransferMode
 import kotlin.native.concurrent.Worker
 
-private data class WorkerArg(val desktopName: String,
-                             val desktop: HDESK?,
+private data class WorkerArg(val desktop: InternalDesktop,
                              val program: String?)
 
+/**
+ * Worker does not allow transferring external references.
+ */
+private data class InternalDesktop(
+    val desktop: HDESK,
+    val name: String?
+)
+
 @Suppress("unused")
 fun main(args: Array<String>?) {
     printf("Secure Desktop demo by i@yuuta.moe")
-    val oldDesktop = GetThreadDesktop(GetCurrentThreadId())
+    val oldDesktop = WinApi.getThreadDesktop()
     val desktopName = "desktop_dialog"
-    val newDesktop = CreateDesktopA(desktopName,
-        null,
-        null,
-        0.convert(),
-        (DESKTOP_READOBJECTS or
+    val newDesktop = WinApi.create(desktopName,
+        DESKTOP_READOBJECTS or
                 DESKTOP_CREATEWINDOW or
                 DESKTOP_CREATEMENU or
                 DESKTOP_HOOKCONTROL or
@@ -30,41 +34,65 @@ fun main(args: Array<String>?) {
                 DESKTOP_JOURNALPLAYBACK or
                 DESKTOP_ENUMERATE or
                 DESKTOP_WRITEOBJECTS or
-                DESKTOP_SWITCHDESKTOP).toUInt(),
-        null)
-    SwitchDesktop(newDesktop)
+                DESKTOP_SWITCHDESKTOP)
+    if (newDesktop == null) {
+        MessageBoxA(null,
+            "Cannot create the desktop. API returned null.",
+            "Secure Desktop",
+            (MB_OK or MB_ICONERROR).toUInt())
+        return
+    }
+    if (oldDesktop == null) {
+        MessageBoxA(null,
+            "Cannot obtain the current desktop. API returned null.",
+            "Secure Desktop",
+            (MB_OK or MB_ICONERROR).toUInt())
+        return
+    }
+    WinApi.switch(newDesktop)
     Worker.start().execute(TransferMode.SAFE,
-        { WorkerArg(desktopName,
-        newDesktop,
+        { WorkerArg(
+            InternalDesktop(newDesktop.desktop,
+                newDesktop.name),
         if (args != null && args.isNotEmpty()) args[0] else null
     ) }) {
         exec(it)
     }.consume {
-        SwitchDesktop(oldDesktop)
-        CloseDesktop(newDesktop)
+        WinApi.switch(oldDesktop)
+        WinApi.close(newDesktop)
     }
 }
 
 private fun exec(arg: WorkerArg): Int {
-    SetThreadDesktop(arg.desktop)
-    if (arg.program != null) memScoped {
-        val startupInfo = cValue<STARTUPINFOA> {
-            lpDesktop = arg.desktopName.cstr.ptr
-            dwFlags = STARTF_RUNFULLSCREEN.convert()
-        }
-        val pi = cValue<PROCESS_INFORMATION> {
+    WinApi.setThreadDesktop(WinDesktop(arg.desktop.desktop,
+        arg.desktop.name))
+    if (arg.program != null) {
+        if (arg.desktop.name != null) {
+            memScoped {
+                val startupInfo = cValue<STARTUPINFOA> {
+                    lpDesktop = arg.desktop.name.cstr.ptr
+                    dwFlags = STARTF_RUNFULLSCREEN.convert()
+                }
+                val pi = cValue<PROCESS_INFORMATION> {
+                }
+                // TODO: Close these objects?
+                return@memScoped CreateProcessA(arg.program,
+                    null,
+                    null,
+                    null,
+                    0,
+                    0.convert(),
+                    null,
+                    null,
+                    startupInfo.ptr,
+                    pi.ptr)
+            }
+        } else {
+            MessageBoxA(null,
+                "Cannot run the program. Cannot obtain the created desktop's name.",
+                "Secure Desktop",
+                (MB_OK or MB_ICONERROR).convert())
         }
-        // TODO: Close these objects?
-        return@memScoped CreateProcessA(arg.program,
-            null,
-            null,
-            null,
-            0,
-            0.convert(),
-            null,
-            null,
-            startupInfo.ptr,
-            pi.ptr)
     }
     return MessageBoxA(null,
         "Hi! This is message from thread ${GetCurrentThreadId()} in the secure desktop." +
diff --git a/src/mingwMain/kotlin/moe.yuuta.desktop/WinApi.kt b/src/mingwMain/kotlin/moe.yuuta.desktop/WinApi.kt
new file mode 100644
index 0000000..2855923
--- /dev/null
+++ b/src/mingwMain/kotlin/moe.yuuta.desktop/WinApi.kt
@@ -0,0 +1,46 @@
+package moe.yuuta.desktop
+
+import kotlinx.cinterop.convert
+import platform.windows.*
+
+object WinApi {
+    /*
+    fun list(lParam: Long, callback: (lpszDesktop: LPSTR?,
+                                      lParam: Long) -> Int): Boolean {
+        return EnumDesktopsA(GetProcessWindowStation(),
+            staticCFunction(callback),
+            lParam) == 1
+    }
+    */
+    fun close(desktop: WinDesktop): Boolean {
+        return CloseDesktop(desktop.desktop) == 1
+    }
+
+    fun getThreadDesktop(thread: Int): WinDesktop? {
+        val desk = GetThreadDesktop(thread.convert()) ?: return null
+        return WinDesktop(desk, null)
+    }
+
+    fun getThreadDesktop(): WinDesktop? = getThreadDesktop(getCurrentThreadId())
+
+    fun getCurrentThreadId(): Int = GetCurrentThreadId().convert()
+
+    fun create(lpszDesktop: String,
+               dwDesiredAccess: Int): WinDesktop? {
+        val desk = CreateDesktopA(lpszDesktop,
+            null,
+            null,
+            0.convert(),
+            dwDesiredAccess.convert(),
+            null) ?: return null
+        return WinDesktop(desk, lpszDesktop)
+    }
+
+    fun switch(desktop: WinDesktop): Boolean {
+        return SwitchDesktop(desktop.desktop) == 1
+    }
+
+    fun setThreadDesktop(desktop: WinDesktop): Boolean {
+        return SetThreadDesktop(desktop.desktop) == 1
+    }
+}
\ No newline at end of file
diff --git a/src/mingwMain/kotlin/moe.yuuta.desktop/WinDesktop.kt b/src/mingwMain/kotlin/moe.yuuta.desktop/WinDesktop.kt
new file mode 100644
index 0000000..4928f9e
--- /dev/null
+++ b/src/mingwMain/kotlin/moe.yuuta.desktop/WinDesktop.kt
@@ -0,0 +1,8 @@
+package moe.yuuta.desktop
+
+import platform.windows.HDESK
+
+data class WinDesktop(
+    val desktop: HDESK,
+    val name: String?
+)
\ No newline at end of file