From 8a9e74e8d7906a6545965cb4e366aa77d9522e14 Mon Sep 17 00:00:00 2001
From: Lau Ching Jun <chingjun@google.com>
Date: Fri, 12 Apr 2024 14:45:32 -0700
Subject: [PATCH] Avoid forwarding the data after socket is disconnected.
 (#146665)

In a ProxiedDevicePortForwarder, there might be a race condition where the local socket has been disconnected, but the remote end was still sending new data. In this case, avoid forwarding new data to the socket.
---
 .../lib/src/proxied_devices/devices.dart      | 13 +++++++++--
 .../proxied_devices/proxied_devices_test.dart | 22 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/packages/flutter_tools/lib/src/proxied_devices/devices.dart b/packages/flutter_tools/lib/src/proxied_devices/devices.dart
index 4b4a9403f49..c286ce2545f 100644
--- a/packages/flutter_tools/lib/src/proxied_devices/devices.dart
+++ b/packages/flutter_tools/lib/src/proxied_devices/devices.dart
@@ -633,11 +633,18 @@ class ProxiedPortForwarder extends DevicePortForwarder {
         'port': devicePort,
       }));
       final Stream<List<int>> dataStream = connection.listenToEvent('proxy.data.$id').asyncExpand((DaemonEventData event) => event.binary);
-      dataStream.listen(socket.add);
+      final StreamSubscription<List<int>> subscription = dataStream.listen(socket.add);
       final Future<DaemonEventData> disconnectFuture = connection.listenToEvent('proxy.disconnected.$id').first;
+
+      bool socketDoneCalled = false;
+
       unawaited(disconnectFuture.then<void>((_) async {
           try {
-            await socket.close();
+            if (socketDoneCalled) {
+              await subscription.cancel();
+            } else {
+              await (subscription.cancel(), socket.close()).wait;
+            }
           } on Exception {
             // ignore
           }
@@ -670,6 +677,8 @@ class ProxiedPortForwarder extends DevicePortForwarder {
         // Do nothing here. Everything will be handled in the `then` block below.
         return false;
       }).whenComplete(() {
+        socketDoneCalled = true;
+        unawaited(subscription.cancel());
         // Send a proxy disconnect event just in case.
         unawaited(connection.sendRequest('proxy.disconnect', <String, Object>{
           'id': id,
diff --git a/packages/flutter_tools/test/general.shard/proxied_devices/proxied_devices_test.dart b/packages/flutter_tools/test/general.shard/proxied_devices/proxied_devices_test.dart
index 4fae8997d87..8efd65c2038 100644
--- a/packages/flutter_tools/test/general.shard/proxied_devices/proxied_devices_test.dart
+++ b/packages/flutter_tools/test/general.shard/proxied_devices/proxied_devices_test.dart
@@ -231,6 +231,28 @@ void main() {
         // Wait the event queue and make sure that it doesn't crash.
         await pumpEventQueue();
       });
+
+      testWithoutContext('should not forward new data to socket after disconnection', () async {
+        // Data will be forwarded before disconnection
+        serverDaemonConnection.sendEvent('proxy.data.$id', null, <int>[1, 2, 3]);
+        await pumpEventQueue();
+        expect(fakeSocket.addedData, <List<int>>[<int>[1, 2, 3]]);
+
+        // It will try to disconnect the remote port when socket is done.
+        fakeSocket.doneCompleter.complete(true);
+        final DaemonMessage message = await broadcastOutput.first;
+
+        expect(message.data['id'], isNotNull);
+        expect(message.data['method'], 'proxy.disconnect');
+        expect(message.data['params'], <String, Object?>{
+          'id': 'random_id',
+        });
+        await pumpEventQueue();
+
+        serverDaemonConnection.sendEvent('proxy.data.$id', null, <int>[4, 5, 6]);
+        await pumpEventQueue();
+        expect(fakeSocket.addedData, <List<int>>[<int>[1, 2, 3]]);
+      });
     });
 
     testWithoutContext('disposes multiple sockets correctly', () async {