Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit d56d7d1

Browse files
ImmutableJeffreyImmutableJeffrey
committedMay 1, 2025·
chore: pkce support poll to event based
1 parent d884787 commit d56d7d1

File tree

2 files changed

+434
-106
lines changed

2 files changed

+434
-106
lines changed
 

‎Source/Immutable/Private/Immutable/Windows/ImmutablePKCEWindows.cpp

Lines changed: 378 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "Immutable/Windows/ImmutablePKCEWindows.h"
22

3+
#include "HAL/Runnable.h"
4+
#include "HAL/RunnableThread.h"
35
#include "Interfaces/IPluginManager.h"
46

57
#include "Windows/AllowWindowsPlatformTypes.h"
@@ -14,12 +16,260 @@ THIRD_PARTY_INCLUDES_END
1416
// Registry constants
1517
static const TCHAR* DeepLinkRegistryValueName = TEXT("DeepLink");
1618

19+
// Static storage for registry watchers by protocol
20+
static TMap<FString, TUniquePtr<FRegistryWatcherRunnable>> GRunnableMap;
21+
22+
// Helper function to get Windows error message from error code
23+
static FString GetWindowsErrorMessage(DWORD ErrorCode)
24+
{
25+
LPWSTR MessageBuffer = nullptr;
26+
DWORD Size = FormatMessageW
27+
(
28+
FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
29+
nullptr, ErrorCode, 0, (LPWSTR)&MessageBuffer, 0, nullptr
30+
);
31+
32+
FString ErrorMessage;
33+
34+
if (MessageBuffer)
35+
{
36+
ErrorMessage = FString(Size, MessageBuffer).TrimEnd();
37+
LocalFree(MessageBuffer);
38+
}
39+
else
40+
{
41+
ErrorMessage = FString::Printf(TEXT("Unknown error code: %d"), ErrorCode);
42+
}
43+
44+
return ErrorMessage;
45+
}
46+
47+
// Helper function to safely clean up a registry watcher thread
48+
static void CleanupThreadForProtocol(const FString& Protocol)
49+
{
50+
if (GRunnableMap.Contains(Protocol))
51+
{
52+
// Remove from map - unique ptr will handle deletion and stopping the thread
53+
GRunnableMap.Remove(Protocol);
54+
}
55+
}
56+
57+
FRegistryWatcherRunnable::FRegistryWatcherRunnable(const FString& InRegistryKeyPath, TWeakObjectPtr<UImmutablePKCEData> InPKCEData):
58+
RegistryKeyPath(InRegistryKeyPath),
59+
PKCEData(InPKCEData)
60+
{
61+
// Create Windows event objects for notification and termination
62+
EventHandle = CreateEvent(NULL, false, false, NULL);
63+
TerminateEventHandle = CreateEvent(NULL, true, false, NULL);
64+
}
65+
66+
FRegistryWatcherRunnable::~FRegistryWatcherRunnable()
67+
{
68+
// Stop the thread if running
69+
FRegistryWatcherRunnable::Stop();
70+
71+
// Wait for thread to complete if it exists
72+
if (Thread)
73+
{
74+
Thread->WaitForCompletion();
75+
Thread.Reset();
76+
}
77+
78+
// Close any open registry key
79+
if (DeepLinkKeyHandle)
80+
{
81+
RegCloseKey(DeepLinkKeyHandle);
82+
DeepLinkKeyHandle = nullptr;
83+
}
84+
85+
// Close event handles
86+
if (EventHandle)
87+
{
88+
CloseHandle(EventHandle);
89+
EventHandle = nullptr;
90+
}
91+
92+
if (TerminateEventHandle)
93+
{
94+
CloseHandle(TerminateEventHandle);
95+
TerminateEventHandle = nullptr;
96+
}
97+
}
98+
99+
bool FRegistryWatcherRunnable::Init()
100+
{
101+
// Validate event handles were created successfully
102+
if (!EventHandle || !TerminateEventHandle)
103+
{
104+
IMTBL_ERR("Failed to create event handles for registry watcher");
105+
return false;
106+
}
107+
108+
return true;
109+
}
110+
111+
bool FRegistryWatcherRunnable::StartThread(const FString& ProtocolName)
112+
{
113+
if (Thread)
114+
{
115+
IMTBL_WARN("Thread already started for registry watcher");
116+
return true;
117+
}
118+
119+
// Create the thread that will run this runnable
120+
Thread = TUniquePtr<FRunnableThread>
121+
(
122+
FRunnableThread::Create
123+
(
124+
this,
125+
*FString::Printf(TEXT("RegistryWatcher_%s"), *ProtocolName),
126+
0, // Use the current thread's stack size
127+
TPri_Normal
128+
)
129+
);
130+
131+
if (!Thread)
132+
{
133+
IMTBL_ERR("Failed to create registry watcher thread for protocol: %s", *ProtocolName);
134+
return false;
135+
}
136+
137+
return true;
138+
}
139+
140+
uint32 FRegistryWatcherRunnable::Run()
141+
{
142+
bool bKeyOpened = false;
143+
144+
HANDLE WaitHandles[2] = {EventHandle, TerminateEventHandle};
145+
146+
// Lambda to close registry key and reset flag
147+
auto CloseRegistryKey = [&]()
148+
{
149+
if (DeepLinkKeyHandle)
150+
{
151+
RegCloseKey(DeepLinkKeyHandle);
152+
DeepLinkKeyHandle = nullptr;
153+
}
154+
bKeyOpened = false;
155+
};
156+
157+
while (bShouldRun)
158+
{
159+
// Open the registry key if not already open
160+
if (!bKeyOpened)
161+
{
162+
// Try to open existing key first
163+
LONG OpenResult = RegOpenKeyExW(HKEY_CURRENT_USER, *RegistryKeyPath, 0, KEY_READ | KEY_WRITE | KEY_NOTIFY, &DeepLinkKeyHandle);
164+
if (OpenResult != ERROR_SUCCESS)
165+
{
166+
// Create the key if it doesn't exist
167+
LONG CreateResult = RegCreateKeyExW(HKEY_CURRENT_USER, *RegistryKeyPath, 0, nullptr, 0, KEY_READ | KEY_WRITE | KEY_NOTIFY, nullptr, &DeepLinkKeyHandle, nullptr);
168+
if (CreateResult != ERROR_SUCCESS)
169+
{
170+
// Failed to open or create the key, log error and retry after delay
171+
IMTBL_ERR("Failed to open or create registry key: %s (Error: %d - %s)", *RegistryKeyPath, CreateResult, *GetWindowsErrorMessage(CreateResult));
172+
FPlatformProcess::Sleep(1.0f);
173+
continue;
174+
}
175+
}
176+
bKeyOpened = true;
177+
178+
// Set up registry change notification
179+
LONG NotifyResult = RegNotifyChangeKeyValue(DeepLinkKeyHandle, false, REG_NOTIFY_CHANGE_LAST_SET, EventHandle, true);
180+
if (NotifyResult != ERROR_SUCCESS)
181+
{
182+
CloseRegistryKey();
183+
FPlatformProcess::Sleep(1.0f);
184+
continue;
185+
}
186+
}
187+
188+
// Wait for the registry change event or termination event
189+
DWORD WaitResult = WaitForMultipleObjects(2, WaitHandles, false, INFINITE);
190+
191+
if (WaitResult == WAIT_OBJECT_0) // Registry change event
192+
{
193+
// Event was signaled, check the registry for deep link
194+
if (DeepLinkKeyHandle && !PKCEData.IsStale() && bShouldRun)
195+
{
196+
// Use Game Thread to process the deeplink
197+
FFunctionGraphTask::CreateAndDispatchWhenReady([this]()
198+
{
199+
if (!PKCEData.IsStale())
200+
{
201+
UImmutablePKCEWindows::CheckDeepLinkRegistry(PKCEData.Get(), DeepLinkKeyHandle);
202+
}
203+
}, TStatId(), nullptr, ENamedThreads::GameThread);
204+
}
205+
206+
// Reset notification for the next change
207+
if (bKeyOpened && bShouldRun)
208+
{
209+
LONG notifyResult = RegNotifyChangeKeyValue(DeepLinkKeyHandle, false, REG_NOTIFY_CHANGE_LAST_SET, EventHandle, true);
210+
if (notifyResult != ERROR_SUCCESS)
211+
{
212+
CloseRegistryKey();
213+
}
214+
}
215+
}
216+
else if (WaitResult == WAIT_OBJECT_0 + 1) // Termination event
217+
{
218+
IMTBL_LOG("Registry watcher received termination signal");
219+
break;
220+
}
221+
else if (WaitResult == WAIT_FAILED)
222+
{
223+
IMTBL_ERR("WaitForMultipleObjects failed with error: %d - %s", GetLastError(), *GetWindowsErrorMessage(GetLastError()));
224+
if (bKeyOpened)
225+
{
226+
CloseRegistryKey();
227+
}
228+
FPlatformProcess::Sleep(0.5f);
229+
}
230+
else
231+
{
232+
// Any other result is unexpected
233+
IMTBL_WARN("Unexpected result from WaitForMultipleObjects: %d", WaitResult);
234+
if (bKeyOpened)
235+
{
236+
CloseRegistryKey();
237+
}
238+
FPlatformProcess::Sleep(0.5f);
239+
}
240+
}
241+
242+
return 0;
243+
}
244+
245+
void FRegistryWatcherRunnable::Stop()
246+
{
247+
bShouldRun = false;
248+
249+
// Signal termination event to wake up the thread
250+
if (TerminateEventHandle)
251+
{
252+
SetEvent(TerminateEventHandle);
253+
}
254+
}
255+
256+
void FRegistryWatcherRunnable::Exit()
257+
{
258+
// Clean up when thread exits
259+
if (DeepLinkKeyHandle)
260+
{
261+
RegCloseKey(DeepLinkKeyHandle);
262+
DeepLinkKeyHandle = nullptr;
263+
}
264+
}
265+
17266
UImmutablePKCEData* UImmutablePKCEWindows::Initialise(const FImmutablePassportInitData& PassportInitData)
18267
{
19268
// Extract the protocol scheme from the redirect URI
20269
FString Protocol;
21270
if (!ExtractProtocolFromUri(PassportInitData.redirectUri, Protocol))
22271
{
272+
IMTBL_ERR("Failed to extract protocol from redirect URI: %s", *PassportInitData.redirectUri);
23273
return nullptr;
24274
}
25275

@@ -35,29 +285,49 @@ UImmutablePKCEData* UImmutablePKCEWindows::Initialise(const FImmutablePassportIn
35285
IMTBL_LOG("Registered protocol handler for %s", *Protocol);
36286

37287
UImmutablePKCEData* PKCEData = NewObject<UImmutablePKCEData>();
38-
39288
PKCEData->PassportInitData = PassportInitData;
40289

41-
// Setup periodic check for deep links
42-
if (GEngine)
290+
// Setup registry watcher thread for deep links
291+
FString RegistryKeyPath = GetDeepLinkRegistryKeyPath(Protocol);
292+
293+
// Clean up any existing thread for this protocol
294+
CleanupThreadForProtocol(Protocol);
295+
296+
// Create new runnable and start its thread
297+
TWeakObjectPtr<UImmutablePKCEData> WeakPKCEData(PKCEData);
298+
TUniquePtr<FRegistryWatcherRunnable> Runnable = MakeUnique<FRegistryWatcherRunnable>(RegistryKeyPath, WeakPKCEData);
299+
300+
if (!Runnable->Init())
43301
{
44-
FTickerDelegate TickDelegate = FTickerDelegate::CreateWeakLambda(PKCEData, [PKCEData](float DeltaTime)
45-
{
46-
return UImmutablePKCEWindows::CheckDeepLinkTick(PKCEData, DeltaTime);
47-
});
48-
PKCEData->TickDelegateHandle = FTSTicker::GetCoreTicker().AddTicker(TickDelegate, 1.0f);
302+
IMTBL_ERR("Failed to initialize registry watcher for protocol: %s", *Protocol);
303+
return PKCEData;
304+
}
305+
306+
if (!Runnable->StartThread(Protocol))
307+
{
308+
IMTBL_ERR("Failed to start registry watcher thread for protocol: %s", *Protocol);
309+
return PKCEData;
49310
}
50311

312+
// Store the runnable in the map
313+
GRunnableMap.Add(Protocol, MoveTemp(Runnable));
314+
51315
return PKCEData;
52316
}
53317

54318
bool UImmutablePKCEWindows::RegisterProtocolHandler(const FString& Protocol)
55319
{
320+
if (Protocol.IsEmpty())
321+
{
322+
IMTBL_ERR("Cannot register protocol handler: Protocol is empty");
323+
return false;
324+
}
325+
56326
// Create a PowerShell script that will handle incoming protocol activations
57327
FString PowerShellFilePath = GetPowerShellScriptPath(Protocol);
58328
FString LogFilePath = GetLogFilePath(Protocol);
59329

60-
auto Cleanup = [&](bool bDeletePowerShellScript, bool bDeleteRegistryKey)
330+
auto CleanupResources = [&](bool bDeletePowerShellScript, bool bDeleteRegistryKey)
61331
{
62332
if (bDeletePowerShellScript)
63333
{
@@ -90,10 +360,11 @@ bool UImmutablePKCEWindows::RegisterProtocolHandler(const FString& Protocol)
90360
// Create main protocol handler registry key
91361
{
92362
HKEY KeyHandle;
93-
if (RegCreateKeyExW(HKEY_CURRENT_USER, *KeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &KeyHandle, nullptr) != ERROR_SUCCESS)
363+
LONG Result = RegCreateKeyExW(HKEY_CURRENT_USER, *KeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &KeyHandle, nullptr);
364+
if (Result != ERROR_SUCCESS)
94365
{
95-
IMTBL_ERR("Failed to create registry key for protocol handler");
96-
Cleanup(true, false);
366+
IMTBL_ERR("Failed to create registry key for protocol handler (Error: %d - %s)", Result, *GetWindowsErrorMessage(Result));
367+
CleanupResources(true, false);
97368
return false;
98369
}
99370

@@ -108,7 +379,7 @@ bool UImmutablePKCEWindows::RegisterProtocolHandler(const FString& Protocol)
108379
if (!bSuccess)
109380
{
110381
IMTBL_ERR("Failed to set registry values for protocol handler");
111-
Cleanup(true, true);
382+
CleanupResources(true, true);
112383
return false;
113384
}
114385
}
@@ -117,22 +388,24 @@ bool UImmutablePKCEWindows::RegisterProtocolHandler(const FString& Protocol)
117388
{
118389
FString ShellKeyPath = FString::Printf(TEXT("%s\\shell"), *KeyPath);
119390
HKEY ShellKeyHandle;
120-
if (RegCreateKeyExW(HKEY_CURRENT_USER, *ShellKeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &ShellKeyHandle, nullptr) != ERROR_SUCCESS)
391+
LONG Result = RegCreateKeyExW(HKEY_CURRENT_USER, *ShellKeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &ShellKeyHandle, nullptr);
392+
if (Result != ERROR_SUCCESS)
121393
{
122-
IMTBL_ERR("Failed to create shell registry key");
123-
Cleanup(true, true);
394+
IMTBL_ERR("Failed to create shell registry key (Error: %d - %s)", Result, *GetWindowsErrorMessage(Result));
395+
CleanupResources(true, true);
124396
return false;
125397
}
126398

127399
FString EmptyValue = TEXT(" "); // Arbitrary value so that it will use the root for the deep link prompt
128-
bool bSuccess = RegSetValueExW(ShellKeyHandle, nullptr, 0, REG_SZ, (const BYTE*)*EmptyValue, (EmptyValue.Len() + 1) * sizeof(TCHAR)) == ERROR_SUCCESS;
400+
LONG SetResult = RegSetValueExW(ShellKeyHandle, nullptr, 0, REG_SZ, (const BYTE*)*EmptyValue, (EmptyValue.Len() + 1) * sizeof(TCHAR));
401+
bool bSuccess = SetResult == ERROR_SUCCESS;
129402

130403
RegCloseKey(ShellKeyHandle);
131404

132405
if (!bSuccess)
133406
{
134-
IMTBL_ERR("Failed to set shell registry value");
135-
Cleanup(true, true);
407+
IMTBL_ERR("Failed to set shell registry value (Error: %d - %s)", SetResult, *GetWindowsErrorMessage(SetResult));
408+
CleanupResources(true, true);
136409
return false;
137410
}
138411
}
@@ -141,91 +414,87 @@ bool UImmutablePKCEWindows::RegisterProtocolHandler(const FString& Protocol)
141414
{
142415
FString CommandKeyPath = FString::Printf(TEXT("%s\\shell\\open\\command"), *KeyPath);
143416
HKEY CommandKeyHandle;
144-
if (RegCreateKeyExW(HKEY_CURRENT_USER, *CommandKeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &CommandKeyHandle, nullptr) != ERROR_SUCCESS)
417+
LONG Result = RegCreateKeyExW(HKEY_CURRENT_USER, *CommandKeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &CommandKeyHandle, nullptr);
418+
if (Result != ERROR_SUCCESS)
145419
{
146-
IMTBL_ERR("Failed to create command registry key");
147-
Cleanup(true, true);
420+
IMTBL_ERR("Failed to create command registry key (Error: %d - %s)", Result, *GetWindowsErrorMessage(Result));
421+
CleanupResources(true, true);
148422
return false;
149423
}
150424

151-
bool bSuccess = RegSetValueExW(CommandKeyHandle, nullptr, 0, REG_SZ, (const BYTE*)*CommandLine, (CommandLine.Len() + 1) * sizeof(TCHAR)) == ERROR_SUCCESS;
425+
LONG SetResult = RegSetValueExW(CommandKeyHandle, nullptr, 0, REG_SZ, (const BYTE*)*CommandLine, (CommandLine.Len() + 1) * sizeof(TCHAR));
426+
bool bSuccess = SetResult == ERROR_SUCCESS;
152427

153428
RegCloseKey(CommandKeyHandle);
154429

155430
if (!bSuccess)
156431
{
157-
IMTBL_ERR("Failed to set command registry value");
158-
Cleanup(true, true);
432+
IMTBL_ERR("Failed to set command registry value (Error: %d - %s)", SetResult, *GetWindowsErrorMessage(SetResult));
433+
CleanupResources(true, true);
159434
return false;
160435
}
161436
}
162437

163438
return true;
164439
}
165440

166-
bool UImmutablePKCEWindows::CheckDeepLinkTick(UImmutablePKCEData* PKCEData, float DeltaTime)
441+
bool UImmutablePKCEWindows::CheckDeepLinkRegistry(UImmutablePKCEData* PKCEData, HKEY DeepLinkKeyHandle)
167442
{
168-
// Extract protocol from the redirect URI
169-
FString Protocol;
170-
if (!ExtractProtocolFromUri(PKCEData->PassportInitData.redirectUri, Protocol))
443+
if (!PKCEData || !DeepLinkKeyHandle)
171444
{
172-
return true; // Continue ticking even if we can't extract the protocol
445+
return false;
173446
}
174447

175-
FString RegistryKeyPath = GetDeepLinkRegistryKeyPath(Protocol);
448+
DWORD ValueType;
449+
DWORD DataSize = 0;
176450

177-
// Check if there's a deep link URL stored in the registry
178-
HKEY DeepLinkKeyHandle;
179-
if (RegOpenKeyExW(HKEY_CURRENT_USER, *RegistryKeyPath, 0, KEY_READ | KEY_WRITE, &DeepLinkKeyHandle) == ERROR_SUCCESS)
451+
// First query to get size of the registry value data
452+
LONG QueryResult = RegQueryValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, nullptr, &ValueType, nullptr, &DataSize);
453+
if (QueryResult == ERROR_SUCCESS && ValueType == REG_SZ && DataSize > 0)
180454
{
181-
DWORD Type;
182-
DWORD DataSize = 0;
455+
TArray<TCHAR> Buffer;
456+
Buffer.SetNumZeroed(DataSize / sizeof(TCHAR) + 1);
183457

184-
// First query to get size of the registry value data
185-
if (RegQueryValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, nullptr, &Type, nullptr, &DataSize) == ERROR_SUCCESS && Type == REG_SZ && DataSize > 0)
458+
// Second query to get actual data content
459+
LONG DataResult = RegQueryValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, nullptr, nullptr, (BYTE*)Buffer.GetData(), &DataSize);
460+
if (DataResult == ERROR_SUCCESS)
186461
{
187-
TArray<TCHAR> Buffer;
188-
Buffer.SetNumZeroed(DataSize / sizeof(TCHAR) + 1);
462+
FString DeepLink = Buffer.GetData();
189463

190-
// Second query to get actual data content
191-
if (RegQueryValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, nullptr, nullptr, (BYTE*)Buffer.GetData(), &DataSize) == ERROR_SUCCESS)
464+
if (!DeepLink.IsEmpty())
192465
{
193-
FString DeepLink = Buffer.GetData();
194-
195-
if (!DeepLink.IsEmpty())
196-
{
197-
IMTBL_LOG("Found deep link in registry: %s", *DeepLink);
466+
IMTBL_LOG("Found deep link in registry: %s", *DeepLink);
198467

199-
// Clear the registry value to avoid processing the same deep link multiple times
200-
FString EmptyValue = TEXT("");
201-
RegSetValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, 0, REG_SZ, (const BYTE*)*EmptyValue, (EmptyValue.Len() + 1) * sizeof(TCHAR));
202-
RegCloseKey(DeepLinkKeyHandle);
468+
// Clear the registry value to avoid processing the same deep link multiple times
469+
FString EmptyValue = TEXT("");
470+
RegSetValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, 0, REG_SZ, (const BYTE*)*EmptyValue, (EmptyValue.Len() + 1) * sizeof(TCHAR));
203471

204-
HandleDeepLink(PKCEData, DeepLink);
205-
206-
return true;
207-
}
472+
HandleDeepLink(PKCEData, DeepLink);
473+
return true;
208474
}
209475
}
210-
211-
RegCloseKey(DeepLinkKeyHandle);
212-
}
213-
else
214-
{
215-
// Create the registry key if it doesn't exist yet
216-
if (RegCreateKeyExW(HKEY_CURRENT_USER, *RegistryKeyPath, 0, nullptr, 0, KEY_WRITE, nullptr, &DeepLinkKeyHandle, nullptr) == ERROR_SUCCESS)
476+
else
217477
{
218-
FString EmptyValue = TEXT("");
219-
RegSetValueExW(DeepLinkKeyHandle, DeepLinkRegistryValueName, 0, REG_SZ, (const BYTE*)*EmptyValue, (EmptyValue.Len() + 1) * sizeof(TCHAR));
220-
RegCloseKey(DeepLinkKeyHandle);
478+
IMTBL_WARN("Failed to read deep link registry value data (Error: %d - %s)", DataResult, *GetWindowsErrorMessage(DataResult));
221479
}
222480
}
481+
else if (QueryResult != ERROR_FILE_NOT_FOUND)
482+
{
483+
// Only log error if it's not due to missing value (which is normal)
484+
IMTBL_WARN("Failed to query deep link registry value size (Error: %d - %s)", QueryResult, *GetWindowsErrorMessage(QueryResult));
485+
}
223486

224-
return true;
487+
return false;
225488
}
226489

227490
void UImmutablePKCEWindows::HandleDeepLink(UImmutablePKCEData* PKCEData, const FString& DeepLink)
228491
{
492+
if (!PKCEData)
493+
{
494+
IMTBL_ERR("Cannot handle deep link: PKCEData is nullptr");
495+
return;
496+
}
497+
229498
// Clean up the deep link URL by removing quotes and extra whitespace
230499
FString CleanDeepLink = DeepLink.TrimStartAndEnd().Replace(TEXT("\""), TEXT(""));
231500
IMTBL_LOG("Handle Deep Link: %s", *CleanDeepLink);
@@ -248,16 +517,20 @@ void UImmutablePKCEWindows::Reset(UImmutablePKCEData* PKCEData)
248517
return;
249518
}
250519

251-
// In shipping builds, clean up protocol handler registry entries and temp files
252-
// In non-shipping builds, skip cleanup for easier testing and debugging
253-
#if UE_BUILD_SHIPPING
254520
// Extract the protocol scheme from the redirect URI
255521
FString Protocol;
256522
if (!ExtractProtocolFromUri(PKCEData->PassportInitData.redirectUri, Protocol))
257523
{
524+
IMTBL_WARN("Failed to extract protocol from URI during reset: %s", *PKCEData->PassportInitData.redirectUri);
258525
return;
259526
}
260527

528+
// Clean up registry watcher thread
529+
CleanupThreadForProtocol(Protocol);
530+
531+
// In shipping builds, clean up protocol handler registry entries and temp files
532+
// In non-shipping builds, skip cleanup for easier testing and debugging
533+
#if UE_BUILD_SHIPPING
261534
// Delete the PowerShell script file
262535
FString PowerShellFilePath = GetPowerShellScriptPath(Protocol);
263536

@@ -289,26 +562,28 @@ void UImmutablePKCEWindows::Reset(UImmutablePKCEData* PKCEData)
289562

290563
// Delete registry keys
291564
FString ProtocolKeyPath = GetProtocolRegistryKeyPath(Protocol);
292-
LSTATUS lResult = RegDeleteTreeW(HKEY_CURRENT_USER, *ProtocolKeyPath);
293-
if (lResult == ERROR_SUCCESS)
565+
LSTATUS DeleteResult = RegDeleteTreeW(HKEY_CURRENT_USER, *ProtocolKeyPath);
566+
if (DeleteResult == ERROR_SUCCESS)
294567
{
295568
IMTBL_LOG("Deleted protocol handler registry key: %s", *ProtocolKeyPath);
296569
}
297-
else
570+
else if (DeleteResult != ERROR_FILE_NOT_FOUND)
298571
{
299-
IMTBL_WARN("Failed to delete protocol handler registry key: %s (Error code: %d)", *ProtocolKeyPath, lResult);
572+
// Only log error if key exists but couldn't be deleted
573+
IMTBL_WARN("Failed to delete protocol handler registry key: %s (Error code: %d - %s)", *ProtocolKeyPath, DeleteResult, *GetWindowsErrorMessage(DeleteResult));
300574
}
301575

302576
// Delete the deep link registry key
303577
FString DeepLinkRegistryKeyPath = GetDeepLinkRegistryKeyPath(Protocol);
304-
lResult = RegDeleteKeyW(HKEY_CURRENT_USER, *DeepLinkRegistryKeyPath);
305-
if (lResult == ERROR_SUCCESS)
578+
DeleteResult = RegDeleteKeyW(HKEY_CURRENT_USER, *DeepLinkRegistryKeyPath);
579+
if (DeleteResult == ERROR_SUCCESS)
306580
{
307581
IMTBL_LOG("Deleted deep link registry key: %s", *DeepLinkRegistryKeyPath);
308582
}
309-
else
583+
else if (DeleteResult != ERROR_FILE_NOT_FOUND)
310584
{
311-
IMTBL_WARN("Failed to delete deep link registry key: %s (Error code: %d)", *DeepLinkRegistryKeyPath, lResult);
585+
// Only log error if key exists but couldn't be deleted
586+
IMTBL_WARN("Failed to delete deep link registry key: %s (Error code: %d - %s)", *DeepLinkRegistryKeyPath, DeleteResult, *GetWindowsErrorMessage(DeleteResult));
312587
}
313588

314589
IMTBL_LOG("PKCE cleanup completed for protocol: %s", *Protocol);
@@ -317,13 +592,6 @@ void UImmutablePKCEWindows::Reset(UImmutablePKCEData* PKCEData)
317592
auto ResetPKCEData = [](UImmutablePKCEData* PKCEData)
318593
{
319594
PKCEData->DynamicMulticastDelegate_DeepLinkCallback.Clear();
320-
321-
if (GEngine)
322-
{
323-
FTSTicker::GetCoreTicker().RemoveTicker(PKCEData->TickDelegateHandle);
324-
PKCEData->TickDelegateHandle.Reset();
325-
}
326-
327595
PKCEData->PassportInitData = {};
328596
};
329597

@@ -332,6 +600,12 @@ void UImmutablePKCEWindows::Reset(UImmutablePKCEData* PKCEData)
332600

333601
FString UImmutablePKCEWindows::GetProtocolDirectory(const FString& Protocol)
334602
{
603+
if (Protocol.IsEmpty())
604+
{
605+
IMTBL_WARN("Empty protocol provided to GetProtocolDirectory, using default");
606+
return FPaths::ProjectSavedDir() / TEXT("Immutable") / TEXT("DefaultProtocol");
607+
}
608+
335609
FString BaseDir;
336610

337611
// In shipping builds, store files in system temp directory
@@ -343,6 +617,7 @@ FString UImmutablePKCEWindows::GetProtocolDirectory(const FString& Protocol)
343617
{
344618
// Fallback if TEMP environment variable is not available
345619
BaseDir = FPaths::ProjectSavedDir();
620+
IMTBL_WARN("System TEMP directory not found, using project saved directory instead");
346621
}
347622
#else
348623
// In non-shipping builds, use saved directory for testing purposes
@@ -357,12 +632,16 @@ FString UImmutablePKCEWindows::GetProtocolDirectory(const FString& Protocol)
357632
#endif
358633

359634
// Ensure the directory exists
360-
if (!FPlatformFileManager::Get().GetPlatformFile().DirectoryExists(*BaseDir))
635+
FString ProtocolDir = FPaths::ConvertRelativePathToFull(BaseDir / Protocol);
636+
if (!FPlatformFileManager::Get().GetPlatformFile().DirectoryExists(*ProtocolDir))
361637
{
362-
FPlatformFileManager::Get().GetPlatformFile().CreateDirectoryTree(*BaseDir);
638+
if (!FPlatformFileManager::Get().GetPlatformFile().CreateDirectoryTree(*ProtocolDir))
639+
{
640+
IMTBL_WARN("Failed to create protocol directory: %s", *ProtocolDir);
641+
}
363642
}
364643

365-
return FPaths::ConvertRelativePathToFull(BaseDir / Protocol);
644+
return ProtocolDir;
366645
}
367646

368647
FString UImmutablePKCEWindows::GetPowerShellScriptPath(const FString& Protocol)
@@ -387,13 +666,24 @@ FString UImmutablePKCEWindows::GetDeepLinkRegistryKeyPath(const FString& Protoco
387666

388667
bool UImmutablePKCEWindows::ExtractProtocolFromUri(const FString& Uri, FString& OutProtocol)
389668
{
390-
FString Right;
669+
if (Uri.IsEmpty())
670+
{
671+
IMTBL_ERR("Cannot extract protocol: URI is empty");
672+
return false;
673+
}
391674

675+
FString Right;
392676
if (!Uri.Split(TEXT("://"), &OutProtocol, &Right))
393677
{
394678
IMTBL_ERR("Failed to extract protocol from URI: %s", *Uri);
395679
return false;
396680
}
397681

682+
if (OutProtocol.IsEmpty())
683+
{
684+
IMTBL_ERR("Extracted protocol is empty from URI: %s", *Uri);
685+
return false;
686+
}
687+
398688
return true;
399689
}

‎Source/Immutable/Public/Immutable/Windows/ImmutablePKCEWindows.h

Lines changed: 56 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,49 @@
22

33
#include "Kismet/BlueprintFunctionLibrary.h"
44

5+
#include "Windows/WindowsHWrapper.h"
6+
57
#include "Immutable/ImmutableDataTypes.h"
68

79
#include "ImmutablePKCEWindows.generated.h"
810

11+
/**
12+
* Self-contained registry watcher that creates and manages its own thread
13+
* Monitors the registry for deep link callbacks and processes them
14+
*/
15+
class FRegistryWatcherRunnable : public FRunnable
16+
{
17+
public:
18+
FRegistryWatcherRunnable(const FString& InRegistryKeyPath, TWeakObjectPtr<UImmutablePKCEData> InPKCEData);
19+
virtual ~FRegistryWatcherRunnable() override;
20+
21+
/** FRunnable: @Interface @Begin */
22+
virtual bool Init() override;
23+
virtual uint32 Run() override;
24+
virtual void Stop() override;
25+
virtual void Exit() override;
26+
/** FRunnable: @Interface @End */
27+
28+
/**
29+
* Starts the watcher thread
30+
* @param ProtocolName Name of the protocol for thread identification
31+
* @return True if thread was started successfully
32+
*/
33+
bool StartThread(const FString& ProtocolName);
34+
35+
private:
36+
FString RegistryKeyPath;
37+
TWeakObjectPtr<UImmutablePKCEData> PKCEData;
38+
volatile bool bShouldRun = true;
39+
40+
HANDLE EventHandle;
41+
HANDLE TerminateEventHandle;
42+
HKEY DeepLinkKeyHandle = nullptr;
43+
44+
// Thread that runs this runnable
45+
TUniquePtr<FRunnableThread> Thread;
46+
};
47+
948
/**
1049
* Windows platform-specific implementation for PKCE authentication flow
1150
* Handles custom URI protocol registration and deep link processing
@@ -30,37 +69,36 @@ class IMMUTABLE_API UImmutablePKCEWindows : public UBlueprintFunctionLibrary
3069
* Creates necessary registry entries and PowerShell helper scripts
3170
*
3271
* @param Protocol The protocol to register (e.g., "immutable")
33-
* @return True if registration was successful
72+
* @return True if registration was successful, false if failed due to invalid protocol or registry errors
3473
*/
3574
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
3675
static bool RegisterProtocolHandler(const FString& Protocol);
3776

3877
/**
39-
* Periodic function to check for deep links in the registry
40-
* Called by the ticker system to detect incoming authentication callbacks
41-
* Monitors registry for URI requests and triggers appropriate delegates
78+
* Process a deep link URL and trigger the appropriate callback
79+
* Parses the incoming URI and extracts authentication parameters
4280
*
43-
* @param PKCEData The PKCE data object containing callback information
44-
* @param DeltaTime Time since last tick
45-
* @return True to continue ticking, false to stop monitoring
81+
* @param PKCEData The PKCE data object containing callback information; if nullptr, function will return early
82+
* @param DeepLink The deep link URL to handle with authorization code and state
4683
*/
4784
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
48-
static bool CheckDeepLinkTick(UImmutablePKCEData* PKCEData, float DeltaTime);
85+
static void HandleDeepLink(UImmutablePKCEData* PKCEData, const FString& DeepLink);
4986

5087
/**
51-
* Process a deep link URL and trigger the appropriate callback
52-
* Parses the incoming URI and extracts authentication parameters
88+
* Check the registry for deep link values using a provided registry key handle
89+
* Called by the registry watcher thread when registry changes are detected
5390
*
5491
* @param PKCEData The PKCE data object containing callback information
55-
* @param DeepLink The deep link URL to handle with authorization code and state
92+
* @param DeepLinkKeyHandle An open registry key handle with appropriate permissions
93+
* @return True if a deep link was found and processed, false if invalid parameters or no deep link found
5694
*/
57-
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
58-
static void HandleDeepLink(UImmutablePKCEData* PKCEData, const FString& DeepLink);
95+
static bool CheckDeepLinkRegistry(UImmutablePKCEData* PKCEData, HKEY DeepLinkKeyHandle);
5996

6097
/**
6198
* Reset and clean up PKCE operation footprint
99+
* Stops any registry watcher threads and cleans up resources
62100
*
63-
* @param PKCEData The PKCE data object to clean up
101+
* @param PKCEData The PKCE data object to clean up; if nullptr, function will return early
64102
*/
65103
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
66104
static void Reset(UImmutablePKCEData* PKCEData);
@@ -71,7 +109,7 @@ class IMMUTABLE_API UImmutablePKCEWindows : public UBlueprintFunctionLibrary
71109
* This script handles deep link registration and browser-to-app communication
72110
*
73111
* @param Protocol The protocol to get the script path for
74-
* @return The full path to the PowerShell script file in the game's saved directory
112+
* @return The full path to the PowerShell script file in the protocol-specific directory
75113
*/
76114
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
77115
static FString GetPowerShellScriptPath(const FString& Protocol);
@@ -81,7 +119,7 @@ class IMMUTABLE_API UImmutablePKCEWindows : public UBlueprintFunctionLibrary
81119
* Used to monitor and debug deep link communication
82120
*
83121
* @param Protocol The protocol to get the log path for
84-
* @return The full path to the log file in the game's saved directory
122+
* @return The full path to the log file in the protocol-specific directory
85123
*/
86124
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
87125
static FString GetLogFilePath(const FString& Protocol);
@@ -92,7 +130,7 @@ class IMMUTABLE_API UImmutablePKCEWindows : public UBlueprintFunctionLibrary
92130
* - In shipping builds: Uses the system temp directory (%TEMP%)
93131
* - In non-shipping builds: Uses a plugin subdirectory in the project's saved folder
94132
*
95-
* @param Protocol The protocol to create a directory for
133+
* @param Protocol The protocol to create a directory for; if empty, returns a default path
96134
* @return The full path to a directory dedicated to the specified protocol
97135
*/
98136
static FString GetProtocolDirectory(const FString& Protocol);
@@ -124,7 +162,7 @@ class IMMUTABLE_API UImmutablePKCEWindows : public UBlueprintFunctionLibrary
124162
*
125163
* @param Uri The URI to extract the protocol from
126164
* @param OutProtocol The output parameter to store the extracted protocol
127-
* @return True if extraction was successful
165+
* @return True if extraction was successful, false if URI is empty or malformed
128166
*/
129167
UFUNCTION(BlueprintCallable, Category = "Immutable|PKCE|Windows")
130168
static bool ExtractProtocolFromUri(const FString& Uri, FString& OutProtocol);

0 commit comments

Comments
 (0)
Please sign in to comment.