Skip to content

Commit

Permalink
Fix shortcut injection bypass with custom workspaces
Browse files Browse the repository at this point in the history
  • Loading branch information
@willtrnr
willtrnr committed May 6, 2021
1 parent ad9fff5 commit 9746760
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 23 deletions.
1 change: 0 additions & 1 deletion AffinityEx.Launcher.Designer/DesignerLauncher.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ public class DesignerLauncher : LauncherBase {
[STAThread]
public static void Main(string[] args) {
Launch(AppName);

}

}
Expand Down
1 change: 0 additions & 1 deletion AffinityEx.Launcher.Photo/PhotoLauncher.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ public class PhotoLauncher : LauncherBase {
[STAThread]
public static void Main(string[] args) {
Launch(AppName);

}

}
Expand Down
1 change: 0 additions & 1 deletion AffinityEx.Launcher.Publisher/PublisherLauncher.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ public class PublisherLauncher : LauncherBase {
[STAThread]
public static void Main(string[] args) {
Launch(AppName);

}

}
Expand Down
52 changes: 34 additions & 18 deletions AffinityEx.Launcher/ApplicationPatches.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,40 +34,44 @@ public static void Apply() {
private static void PatchApplication(Harmony harmony) {
harmony.Patch(
original: AccessTools.PropertyGetter(typeof(Serif.Interop.Persona.Application), "InstallationDirectory"),
prefix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_InstallationDirectory_Prefix))
prefix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_InstallationDirectory_Prefix))
);
harmony.Patch(
original: AccessTools.Method(typeof(Serif.Affinity.Application), "OnStartup"),
prefix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_OnStartup_Prefix))
prefix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_OnStartup_Prefix))
);
harmony.Patch(
original: AccessTools.Method(typeof(Serif.Affinity.Application), "InitialiseServices"),
postfix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_InitialiseServices_Postfix))
postfix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_InitialiseServices_Postfix))
);
harmony.Patch(
original: AccessTools.Method(typeof(Serif.Affinity.Application), "OnServicesInitialised"),
postfix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_OnServicesInitialised_Postfix))
postfix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_OnServicesInitialised_Postfix))
);
harmony.Patch(
original: AccessTools.Method(typeof(Serif.Affinity.Application), "OnMainWindowLoaded"),
postfix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_OnMainWindowLoaded_Postfix))
postfix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_OnMainWindowLoaded_Postfix))
);
harmony.Patch(
original: AccessTools.Method(typeof(Serif.Affinity.Application), "OnFirstIdle"),
postfix: new HarmonyMethod(typeof(Impl), nameof(Impl.Application_OnFirstIdle_Postfix))
postfix: new HarmonyMethod(typeof(Patched), nameof(Patched.Application_OnFirstIdle_Postfix))
);
}

private static void PatchWorkspaces(Harmony harmony) {
var menuPatch = new HarmonyMethod(typeof(Impl), nameof(Impl.Workspace_GetDefaultMenu_Postfix));
var shortcutsPatch = new HarmonyMethod(typeof(Impl), nameof(Impl.Workspace_GetDefaultShortcuts_Postfix));
var menuPatch = new HarmonyMethod(typeof(Patched), nameof(Patched.Workspace_GetDefaultMenu_Postfix));
var defaultShortcutsPatch = new HarmonyMethod(typeof(Patched), nameof(Patched.Workspace_GetDefaultShortcuts_Postfix));
foreach (Type type in workspaceTypes) {
harmony.Patch(AccessTools.Method(type, "GetDefaultMenu"), postfix: menuPatch);
harmony.Patch(AccessTools.Method(type, "GetDefaultShortcuts"), postfix: shortcutsPatch);
harmony.Patch(AccessTools.Method(type, "GetDefaultShortcuts"), postfix: defaultShortcutsPatch);
}
harmony.Patch(
AccessTools.Method(typeof(Workspace), "RemoveInvalidShortcuts"),
postfix: new HarmonyMethod(typeof(Patched), nameof(Patched.Workspace_RemoveInvalidShortcuts_Postfix))
);
}

private static class Impl {
private static class Patched {

internal static bool Application_InstallationDirectory_Prefix(ref string __result) {
__result = AppContext.Current.InstallationDirectory;
Expand Down Expand Up @@ -127,19 +131,31 @@ internal static void Workspace_GetDefaultMenu_Postfix(Workspace __instance, ref
}
}

internal static void Workspace_GetDefaultShortcuts_Postfix(Workspace __instance, ref WorkspaceShortcuts __result) {
internal static void Workspace_GetDefaultShortcuts_Postfix(Workspace __instance, WorkspaceShortcuts __result) {
Log.Debug("Intercepted GetDefaultShortcuts for workspace {WorkspaceName}, forwarding to plugins", __instance.Name);
InjectPluginShortcuts(__instance, __result);
}

internal static void Workspace_RemoveInvalidShortcuts_Postfix(Workspace __instance, WorkspaceShortcuts ___m_shortcuts) {
Log.Debug("Intercepted Shortcuts setter for workspace {WorkspaceName}, forwarding to plugins", __instance.Name);
InjectPluginShortcuts(__instance, ___m_shortcuts);
}

private static void InjectPluginShortcuts(Workspace workspace, WorkspaceShortcuts result) {
var injected = false;
foreach (var plugin in AppContext.Current.Plugins) {
var pluginShortcuts = plugin.GetShortcuts(__instance);
var pluginShortcuts = plugin.GetShortcuts(workspace);
if (pluginShortcuts != null) {
__result.Commands.AddRange(pluginShortcuts.Commands);
__result.GlobalCommands.AddRange(pluginShortcuts.GlobalCommands);
__result.ToolTypes.AddRange(pluginShortcuts.ToolTypes);
__result.ToolKeys.AddRange(pluginShortcuts.ToolKeys);
result.Commands.AddRange(pluginShortcuts.Commands);
result.GlobalCommands.AddRange(pluginShortcuts.GlobalCommands);
result.ToolTypes.AddRange(pluginShortcuts.ToolTypes);
result.ToolKeys.AddRange(pluginShortcuts.ToolKeys);
injected = true;
}
}
// Bindings are cached on first read, we need to blow this value to regenerate bindings for our new commands
__result.Bindings = null;
if (injected) {
result.Bindings = null;
}
}

}
Expand Down
4 changes: 2 additions & 2 deletions AffinityEx.Launcher/NetRuntimePatches.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@ public static void Apply() {
private static void PatchEntryAssembly(Harmony harmony) {
harmony.Patch(
original: AccessTools.Method(typeof(Assembly), "GetEntryAssembly"),
prefix: new HarmonyMethod(typeof(Impl), nameof(Impl.Assembly_GetEntryAssembly_Prefix))
prefix: new HarmonyMethod(typeof(Patched), nameof(Patched.Assembly_GetEntryAssembly_Prefix))
);
}

private static class Impl {
private static class Patched {

internal static bool Assembly_GetEntryAssembly_Prefix(ref Assembly __result) {
var ctx = AppContext.Current;
Expand Down

0 comments on commit 9746760

Please sign in to comment.