1515#import < objc/runtime.h>
1616#import < objc/message.h>
1717
18+ static dispatch_once_t injectMenuOnceToken;
19+
1820WindowTransparencyController *transparencyController;
1921NSMenu *AfloatXMenu;
2022NSMenuItem *AfloatXItem;
3032NSMenu *windowOutlineSubmenu;
3133NSArray *afloatXItems;
3234CIFilter * colorInvertFilter;
33- BOOL menuInjected;
3435
3536@interface AfloatX ()
3637
@@ -165,9 +166,7 @@ + (void)load {
165166
166167 colorInvertFilter = [CIFilter filterWithName: @" CIColorInvert" ];
167168 [colorInvertFilter setDefaults ];
168-
169- menuInjected = NO ;
170-
169+
171170 transparencyController = [WindowTransparencyController sharedInstance ];
172171
173172 AfloatXMenu = [NSMenu new ];
@@ -212,57 +211,71 @@ + (void)load {
212211 transparencyItem,
213212 nil ];
214213 [AfloatXSubmenu setItemArray: afloatXItems];
215-
216- // If the application has a custom dock menu, we will add ourselves to that
217- if ([[NSApp delegate ] respondsToSelector: @selector (applicationDockMenu: )]) {
218- AfloatXMenu = [[NSApp delegate ] applicationDockMenu: NSApp ];
219- [AfloatXMenu addItem: [NSMenuItem separatorItem ]];
220- menuInjected = YES ;
214+
215+ // Are we in an Electron app?
216+ if (objc_lookUpClass (" AtomApplicationDelegate" )) {
217+ dispatch_once (&injectMenuOnceToken, ^{ /* Use up token */ });
218+ [AfloatXMenu addItem: AfloatXItem];
221219 }
222-
223- [AfloatXMenu addItem: AfloatXItem];
224- _ZKSwizzle ([AXAppDelegate class ], [[NSApp delegate ] class ]);
225220}
226221
227222@end
228223
224+ /*
225+ AtomApplicationDelegate is the application delegate class
226+ for all Electron apps. In order for AfloatX to work we must
227+ explicitly swizzle this class because the application delegate
228+ class gets subclassed/swizzled to this one after injection.
229+ */
230+ ZKSwizzleInterface (AXAppDelegate, AtomApplicationDelegate, NSObject )
229231@implementation AXAppDelegate
230232- (NSMenu *)applicationDockMenu:(NSApplication *)sender {
231- if (menuInjected) {
233+ NSMenu *originalMenu = ZKOrig (NSMenu *, sender);
234+ if (originalMenu) {
232235 [AfloatXMenu removeItem: AfloatXItem];
233- AfloatXMenu = ZKOrig ( NSMenu *, sender) ;
236+ AfloatXMenu = originalMenu ;
234237 // Only add a separator if last item isn't already a separator
235238 if (!AfloatXMenu.itemArray .lastObject .isSeparatorItem )
236239 [AfloatXMenu addItem: [NSMenuItem separatorItem ]];
237240 [AfloatXMenu addItem: AfloatXItem];
238241 }
242+
239243 return AfloatXMenu;
240244}
241245@end
242246
243247ZKSwizzleInterface (AXApplication, NSApplication , NSResponder )
244248@implementation AXApplication
245- - (CFArrayRef)_flattenMenu:(NSMenu *)arg1 flatList:(NSArray *)arg2 {
249+ - (CFArrayRef)_flattenMenu:(NSMenu *)dockMenu flatList:(NSArray *)flatList {
250+
251+ // Add AfloatX to the dock menu
252+ dispatch_once (&injectMenuOnceToken, ^{
253+ if (!dockMenu.itemArray .lastObject .isSeparatorItem )
254+ [dockMenu addItem: [NSMenuItem separatorItem ]];
255+
256+ [dockMenu addItem: AfloatXItem];
257+ });
258+
246259 // Make any necessary changes to our menu before it is 'flattened'
247260 NSWindow *window = [AXWindowUtils windowToModify ];
248261 if (!window) {
249262 AfloatXItem.enabled = NO ;
250- return ZKOrig (CFArrayRef, arg1, arg2 );
263+ return ZKOrig (CFArrayRef, dockMenu, flatList );
251264 }
252-
265+
253266 AfloatXItem.enabled = YES ;
254267 if ([[AfloatX sharedInstance ] isWindowTransient: window]) {
255268 [transientItem setState: NSControlStateValueOn ];
256269 } else {
257270 [transientItem setState: NSControlStateValueOff ];
258271 }
259-
272+
260273 if ([[AfloatX sharedInstance ] isWindowSticky: window]) {
261274 [stickyItem setState: NSControlStateValueOn ];
262275 } else {
263276 [stickyItem setState: NSControlStateValueOff ];
264277 }
265-
278+
266279 CGWindowLevel windowLevel = [AXWindowUtils getCGWindowLevelForWindow: window];
267280 if (windowLevel != kCGNormalWindowLevel ) {
268281 if (windowLevel == kCGBackstopMenuLevel ) {
@@ -276,13 +289,13 @@ - (CFArrayRef)_flattenMenu:(NSMenu *)arg1 flatList:(NSArray *)arg2 {
276289 [dropItem setState: NSControlStateValueOff ];
277290 [floatItem setState: NSControlStateValueOff ];
278291 }
279-
292+
280293 if ([AXWindowUtils window: window hasLowTag: CGSTagTransparent]) {
281294 [clickPassthroughItem setState: NSControlStateValueOn ];
282295 } else {
283296 [clickPassthroughItem setState: NSControlStateValueOff ];
284297 }
285-
298+
286299 /* Create a new WindowOutliningController per window */
287300 if (!objc_getAssociatedObject (window, " outlineController" )) {
288301 WindowOutliningController *outlineController = [WindowOutliningController new ];
@@ -292,13 +305,13 @@ - (CFArrayRef)_flattenMenu:(NSMenu *)arg1 flatList:(NSArray *)arg2 {
292305 WindowOutliningController *outlineController = objc_getAssociatedObject (window, " outlineController" );
293306 windowOutlineSubmenu.itemArray = [outlineController colorItems ];
294307 }
295-
308+
296309 if ([objc_getAssociatedObject (window, " isColorInverted" ) boolValue ]) {
297310 [invertColorItem setState: NSControlStateValueOn ];
298311 } else {
299312 [invertColorItem setState: NSControlStateValueOff ];
300313 }
301-
302- return ZKOrig (CFArrayRef, arg1, arg2 );
314+
315+ return ZKOrig (CFArrayRef, dockMenu, flatList );
303316}
304317@end
0 commit comments