fix: Avoid keeping strong reference to self instance in delegates

WebDriverAgentLib/Routing/FBTCPSocket.h

@@ -38,7 +38,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FBTCPSocket : NSObject
 
-@property (nullable, nonatomic) id<FBTCPSocketDelegate> delegate;
+#if __has_feature(objc_arc_weak)
+@property (nullable, nonatomic, weak) id<FBTCPSocketDelegate> delegate;
+#else
+@property (nullable, nonatomic, assign) id<FBTCPSocketDelegate> delegate;
+#endif
 
 /**
  Creates TCP socket isntance which is going to be started on the specified port

WebDriverAgentLib/Routing/FBTCPSocket.m

@@ -48,11 +48,14 @@ - (BOOL)startWithError:(NSError **)error
 - (void)stop
 {
   @synchronized(self.connectedClients) {
-    for (NSUInteger i = 0; i < [self.connectedClients count]; i++) {
-      [[self.connectedClients objectAtIndex:i] disconnect];
+    NSArray *clients = self.connectedClients.copy;
+    [self.connectedClients removeAllObjects];
+    for (GCDAsyncSocket *client in clients) {
+      [client disconnect];
     }
   }
 
+  self.delegate = nil;
   [self.listeningSocket disconnect];
 }
 
@@ -66,20 +69,29 @@ - (void)socket:(GCDAsyncSocket *)sock didAcceptNewSocket:(GCDAsyncSocket *)newSo
   @synchronized(self.connectedClients) {
     [self.connectedClients addObject:newSocket];
   }
-  [self.delegate didClientConnect:newSocket];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientConnect:newSocket];
+  }
 }
 
 - (void)socket:(GCDAsyncSocket *)sock didReadData:(NSData *)data withTag:(long)tag
 {
-  [self.delegate didClientSendData:sock];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientSendData:sock];
+  }
 }
 
 - (void)socketDidDisconnect:(GCDAsyncSocket *)sock withError:(NSError *)err
 {
   @synchronized(self.connectedClients) {
     [self.connectedClients removeObject:sock];
   }
-  [self.delegate didClientDisconnect:sock];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientDisconnect:sock];
+  }
 }
 
 @end

WebDriverAgentLib/Routing/FBWebServer.m

@@ -47,10 +47,16 @@ @interface FBWebServer ()
 @property (nonatomic, strong) RoutingHTTPServer *server;
 @property (atomic, assign) BOOL keepAlive;
 @property (nonatomic, nullable) FBTCPSocket *screenshotsBroadcaster;
+@property (nonatomic, nullable, strong) FBMjpegServer *mjpegServer;
 @end
 
 @implementation FBWebServer
 
+- (void)dealloc
+{
+  [self stopScreenshotsBroadcaster];
+}
+
 + (NSArray<Class<FBCommandHandler>> *)collectCommandHandlerClasses
 {
   NSArray *handlersClasses = FBClassesThatConformsToProtocol(@protocol(FBCommandHandler));
@@ -97,7 +103,7 @@ - (void)startHTTPServer
     [self.server setInterface:bindingIP];
     [FBLogger logFmt:@"Using custom binding IP address: %@", bindingIP];
   }
-  
+
   NSError *error;
   BOOL serverStarted = NO;
 
@@ -117,31 +123,42 @@ - (void)startHTTPServer
     [FBLogger logFmt:@"Last attempt to start web server failed with error %@", [error description]];
     abort();
   }
-  
+
   NSString *serverHost = bindingIP ?: ([XCUIDevice sharedDevice].fb_wifiIPAddress ?: @"127.0.0.1");
   [FBLogger logFmt:@"%@http://%@:%d%@", FBServerURLBeginMarker, serverHost, [self.server port], FBServerURLEndMarker];
 }
 
 - (void)initScreenshotsBroadcaster
 {
   [self readMjpegSettingsFromEnv];
+  self.mjpegServer = [[FBMjpegServer alloc] init];
   self.screenshotsBroadcaster = [[FBTCPSocket alloc]
                                  initWithPort:(uint16_t)FBConfiguration.mjpegServerPort];
-  self.screenshotsBroadcaster.delegate = [[FBMjpegServer alloc] init];
+  self.screenshotsBroadcaster.delegate = self.mjpegServer;
   NSError *error;
   if (![self.screenshotsBroadcaster startWithError:&error]) {
     [FBLogger logFmt:@"Cannot init screenshots broadcaster service on port %@. Original error: %@", @(FBConfiguration.mjpegServerPort), error.description];
+    [self.mjpegServer stopStreaming];
+    self.mjpegServer = nil;
     self.screenshotsBroadcaster = nil;
   }
 }
 
 - (void)stopScreenshotsBroadcaster
 {
   if (nil == self.screenshotsBroadcaster) {
+    self.mjpegServer = nil;
     return;
   }
 
+  id<FBTCPSocketDelegate> delegate = self.screenshotsBroadcaster.delegate;
+  if ([(NSObject *)delegate respondsToSelector:@selector(stopStreaming)]) {
+    [(FBMjpegServer *)delegate stopStreaming];
+  }
+  self.screenshotsBroadcaster.delegate = nil;
   [self.screenshotsBroadcaster stop];
+  self.screenshotsBroadcaster = nil;
+  self.mjpegServer = nil;
 }
 
 - (void)readMjpegSettingsFromEnv
@@ -164,6 +181,8 @@ - (void)stopServing
   if (self.server.isRunning) {
     [self.server stop:NO];
   }
+  self.server = nil;
+  self.exceptionHandler = nil;
   self.keepAlive = NO;
 }
 
@@ -192,10 +211,15 @@ - (BOOL)attemptToStartServer:(RoutingHTTPServer *)server onPort:(NSInteger)port
 
 - (void)registerRouteHandlers:(NSArray *)commandHandlerClasses
 {
+  __weak typeof(self) weakSelf = self;
   for (Class<FBCommandHandler> commandHandler in commandHandlerClasses) {
     NSArray *routes = [commandHandler routes];
     for (FBRoute *route in routes) {
       [self.server handleMethod:route.verb withPath:route.path block:^(RouteRequest *request, RouteResponse *response) {
+        __strong typeof(weakSelf) strongSelf = weakSelf;
+        if (nil == strongSelf) {
+          return;
+        }
         NSDictionary *arguments = [NSJSONSerialization JSONObjectWithData:request.body options:NSJSONReadingMutableContainers error:NULL];
         FBRouteRequest *routeParams = [FBRouteRequest
           routeRequestWithURL:request.url
@@ -209,7 +233,7 @@ - (void)registerRouteHandlers:(NSArray *)commandHandlerClasses
           [route mountRequest:routeParams intoResponse:response];
         }
         @catch (NSException *exception) {
-          [self handleException:exception forResponse:response];
+          [strongSelf handleException:exception forResponse:response];
         }
       }];
     }
@@ -237,9 +261,14 @@ - (void)registerServerKeyRouteHandlers
     [response respondWithString:calibrationPage];
   }];
 
+  __weak typeof(self) weakSelf = self;
   [self.server get:@"/wda/shutdown" withBlock:^(RouteRequest *request, RouteResponse *response) {
+    __strong typeof(weakSelf) strongSelf = weakSelf;
+    if (nil == strongSelf) {
+      return;
+    }
     [response respondWithString:@"Shutting down"];
-    [self.delegate webServerDidRequestShutdown:self];
+    [strongSelf.delegate webServerDidRequestShutdown:strongSelf];
   }];
 
   [self registerRouteHandlers:@[FBUnknownCommands.class]];

WebDriverAgentLib/Utilities/FBImageProcessor.m

@@ -27,6 +27,7 @@ @interface FBImageProcessor ()
 @property (nonatomic) NSData *nextImage;
 @property (nonatomic, readonly) NSLock *nextImageLock;
 @property (nonatomic, readonly) dispatch_queue_t scalingQueue;
+@property (atomic, assign) BOOL isScalingScheduled;
 
 @end
 
@@ -38,6 +39,7 @@ - (id)init
   if (self) {
     _nextImageLock = [[NSLock alloc] init];
     _scalingQueue = dispatch_queue_create("image.scaling.queue", NULL);
+    _isScalingScheduled = NO;
   }
   return self;
 }
@@ -51,32 +53,45 @@ - (void)submitImageData:(NSData *)image
     [FBLogger verboseLog:@"Discarding screenshot"];
   }
   self.nextImage = image;
+  BOOL shouldSchedule = !self.isScalingScheduled;
+  if (shouldSchedule) {
+    self.isScalingScheduled = YES;
+  }
   [self.nextImageLock unlock];
+  if (!shouldSchedule) {
+    return;
+  }
 
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wcompletion-handler"
   dispatch_async(self.scalingQueue, ^{
-    [self.nextImageLock lock];
-    NSData *nextImageData = self.nextImage;
-    self.nextImage = nil;
-    [self.nextImageLock unlock];
-    if (nextImageData == nil) {
-      return;
-    }
+    while (YES) {
+      @autoreleasepool {
+        [self.nextImageLock lock];
+        NSData *nextImageData = self.nextImage;
+        self.nextImage = nil;
+        if (nextImageData == nil) {
+          self.isScalingScheduled = NO;
+          [self.nextImageLock unlock];
+          return;
+        }
+        [self.nextImageLock unlock];
 
-    // We do not want this value to be too high because then we get images larger in size than original ones
-    // Although, we also don't want to lose too much of the quality on recompression
-    CGFloat recompressionQuality = MAX(0.9,
-                                       MIN(FBMaxCompressionQuality, FBConfiguration.mjpegServerScreenshotQuality / 100.0));
-    NSData *thumbnailData = [self.class fixedImageDataWithImageData:nextImageData
-                                                      scalingFactor:scalingFactor
-                                                                uti:UTTypeJPEG
-                                                 compressionQuality:recompressionQuality
-    // iOS always returns screnshots in portrait orientation, but puts the real value into the metadata
-    // Use it with care. See https://github.com/appium/WebDriverAgent/pull/812
-                                                     fixOrientation:FBConfiguration.mjpegShouldFixOrientation
-                                                 desiredOrientation:nil];
-    completionHandler(thumbnailData ?: nextImageData);
+        // We do not want this value to be too high because then we get images larger in size than original ones
+        // Although, we also don't want to lose too much of the quality on recompression
+        CGFloat recompressionQuality = MAX(0.9,
+                                           MIN(FBMaxCompressionQuality, FBConfiguration.mjpegServerScreenshotQuality / 100.0));
+        NSData *thumbnailData = [self.class fixedImageDataWithImageData:nextImageData
+                                                          scalingFactor:scalingFactor
+                                                                    uti:UTTypeJPEG
+                                                     compressionQuality:recompressionQuality
+        // iOS always returns screenshots in portrait orientation, but puts the real value into the metadata
+        // Use it with care. See https://github.com/appium/WebDriverAgent/pull/812
+                                                         fixOrientation:FBConfiguration.mjpegShouldFixOrientation
+                                                     desiredOrientation:nil];
+        completionHandler(thumbnailData ?: nextImageData);
+      }
+    }
   });
 #pragma clang diagnostic pop
 }

WebDriverAgentLib/Utilities/FBMjpegServer.h

@@ -19,6 +19,11 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (instancetype)init;
 
+/**
+ Stops screenshot broadcasting and prevents future scheduling.
+ */
+- (void)stopStreaming;
+
 @end
 
 NS_ASSUME_NONNULL_END