fix: correctly cleaning up the websocket streams

aws-amplify · Dec 19, 2024 · e3729cf · e3729cf
1 parent a15f3f5
commit e3729cf
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 8 deletions.
diff --git a/AWSIoT/Internal/AWSIoTMQTTClient.m b/AWSIoT/Internal/AWSIoTMQTTClient.m
@@ -705,7 +705,12 @@ - (void)invalidateReconnectTimer {
 
 - (void)cleanUpWebsocketOutputStream {
     @synchronized(self) {
-        if (self.websocketOutputStream) {
+        // Before cleaning up the websocket output stream, we must apply stricter stream status checks to avoid possible cocurrent access issue, because the same stream object is possible to be shared in multiple one thread, as for this stream, e.g. the `AWSIoTStreamThread`.
+        if (
+            self.websocketOutputStream && 
+            self.websocketOutputStream.delegate && 
+            (self.websocketOutputStream.streamStatus != NSStreamStatusNotOpen && self.websocketOutputStream.streamStatus != NSStreamStatusClosed)
+        ) {
             self.websocketOutputStream.delegate = nil;
             [self.websocketOutputStream close];
             [self.websocketOutputStream removeFromRunLoop:[NSRunLoop currentRunLoop] forMode:NSDefaultRunLoopMode];
@@ -1258,7 +1263,17 @@ - (void)webSocket:(AWSSRWebSocket *)webSocket didFailWithError:(NSError *)error
     // Also, the webSocket can be set to nil
     [self cleanUpWebsocketOutputStream];
 
-    [self.encoderOutputStream close];
+    // Before cleaning up the stream, we must apply stricter stream status checks to avoid possible cocurrent access issue, because the same stream object is possible to be shared in multiple threads.
+    if (
+        self.encoderOutputStream && 
+        self.encoderOutputStream.delegate && 
+        (self.encoderOutputStream.streamStatus != NSStreamStatusNotOpen && self.encoderOutputStream.streamStatus != NSStreamStatusClosed)
+    ) {
+        self.encoderOutputStream.delegate = nil;
+        [self.encoderOutputStream close];
+        self.encoderOutputStream = nil;
+    }
+
     [self.webSocket close];
     self.webSocket = nil;
 
@@ -1296,7 +1311,17 @@ - (void)webSocket:(AWSSRWebSocket *)webSocket didCloseWithCode:(NSInteger)code r
     // The WebSocket has closed. The input/output streams can be closed here.
     [self cleanUpWebsocketOutputStream];
 
-    [self.encoderOutputStream close];
+    // Before cleaning up the stream, we must apply stricter stream status checks to avoid possible cocurrent access issue, because the same stream object is possible to be shared in multiple threads.
+    if (
+        self.encoderOutputStream && 
+        self.encoderOutputStream.delegate && 
+        (self.encoderOutputStream.streamStatus != NSStreamStatusNotOpen && self.encoderOutputStream.streamStatus != NSStreamStatusClosed)
+    ) {
+        self.encoderOutputStream.delegate = nil;
+        [self.encoderOutputStream close];
+        self.encoderOutputStream = nil;
+    }
+
     [self.webSocket close];
     self.webSocket = nil;
 

diff --git a/AWSIoT/Internal/AWSIoTStreamThread.m b/AWSIoT/Internal/AWSIoTStreamThread.m
@@ -148,7 +148,12 @@ - (void)cleanUp {
                 self.session = nil;
             }
 
-            if (self.outputStream) {
+             // Before cleaning up the output stream, we must apply stricter stream status checks to avoid possible cocurrent access issue, because the same stream object is possible to be shared in mutliple threads.
+            if (
+                self.outputStream && 
+                self.outputStream.delegate && 
+                (self.outputStream.streamStatus != NSStreamStatusNotOpen && self.outputStream.streamStatus != NSStreamStatusClosed)
+            ) {
                 self.outputStream.delegate = nil;
                 [self.outputStream close];
                 [self.outputStream removeFromRunLoop:self.runLoopForStreamsThread

diff --git a/AWSIoT/Internal/MQTTSDK/AWSMQTTEncoder.m b/AWSIoT/Internal/MQTTSDK/AWSMQTTEncoder.m
@@ -44,10 +44,13 @@ - (void)open {
 }
 
 - (void)close {
-    AWSDDLogDebug(@"closing encoder stream.");
-    [self.stream close];
-    [self.stream setDelegate:nil];
-    self.stream = nil;
+    // Before cleaning up the stream, we must apply stricter stream status checks to avoid possible cocurrent access issue, because the same stream object is possible to be shared in multiple threads.
+    if (self.stream && self.stream.delegate && (self.stream.streamStatus != NSStreamStatusNotOpen && self.stream.streamStatus != NSStreamStatusClosed)) {
+        AWSDDLogDebug(@"closing encoder stream.");
+        self.stream.delegate = nil;
+        [self.stream close];
+        self.stream = nil;
+    }
 }
 
 //This is executed in the runLoop.