1 files changed, 61 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." +