From a3bff77649df22fc99de0118fe4bd9893a2318d8 Mon Sep 17 00:00:00 2001
From: pedro ernesto <pedroernestovogado@gmail.com>
Date: Tue, 9 Jun 2026 16:18:54 -0300
Subject: [PATCH] fix: prevent stream localId collisions from destroying live
 stream state

AdbConnection allocated stream IDs with random.nextInt() and never checked
them against streams already open on the connection. MessageQueue keys all
stream state on that raw ID, and stopListening removes it, so when two
concurrently-open streams collided on an ID, the first one to close wiped
the survivor's queues entry and its next take() failed with the fatal
IllegalStateException("Not listening for localId") even though it was
never closed.

- Allocate IDs sequentially (AtomicInteger), like AOSP's adb client, so
  collisions within a connection are impossible.
- startListening now fails fast on a duplicate ID instead of silently
  sharing state between two streams.
- A missing queues entry in poll() now surfaces as the graceful
  AdbStreamClosed (which callers already handle) instead of
  IllegalStateException.
---
 dadb/src/main/kotlin/dadb/AdbConnection.kt    |  9 +++++---
 dadb/src/main/kotlin/dadb/MessageQueue.kt     |  4 ++--
 dadb/src/test/kotlin/dadb/MessageQueueTest.kt | 23 +++++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/dadb/src/main/kotlin/dadb/AdbConnection.kt b/dadb/src/main/kotlin/dadb/AdbConnection.kt
index 57cd17e..8322e46 100644
--- a/dadb/src/main/kotlin/dadb/AdbConnection.kt
+++ b/dadb/src/main/kotlin/dadb/AdbConnection.kt
@@ -25,7 +25,7 @@ import org.jetbrains.annotations.TestOnly
 import java.io.Closeable
 import java.io.IOException
 import java.net.Socket
-import java.util.*
+import java.util.concurrent.atomic.AtomicInteger
 
 internal class AdbConnection internal constructor(
         adbReader: AdbReader,
@@ -36,7 +36,7 @@ internal class AdbConnection internal constructor(
         private val maxPayloadSize: Int
 ) : AutoCloseable {
 
-    private val random = Random()
+    private val nextLocalId = AtomicInteger(0)
     private val messageQueue = AdbMessageQueue(adbReader)
 
     @Throws(IOException::class)
@@ -58,8 +58,11 @@ internal class AdbConnection internal constructor(
         return supportedFeatures.contains(feature)
     }
 
+    // Sequential like AOSP's adb client. Random ints can collide with an ID
+    // already in use; the colliding stream's close then destroys the live
+    // stream's message queue ("Not listening for localId").
     private fun newId(): Int {
-        return random.nextInt()
+        return nextLocalId.incrementAndGet()
     }
 
     @TestOnly
diff --git a/dadb/src/main/kotlin/dadb/MessageQueue.kt b/dadb/src/main/kotlin/dadb/MessageQueue.kt
index b9e7b47..37a50ee 100644
--- a/dadb/src/main/kotlin/dadb/MessageQueue.kt
+++ b/dadb/src/main/kotlin/dadb/MessageQueue.kt
@@ -46,7 +46,7 @@ internal abstract class MessageQueue<V> {
     }
 
     fun startListening(localId: Int) {
-        openStreams.add(localId)
+        check(openStreams.add(localId)) { "Already listening for localId: $localId" }
         queues.putIfAbsent(localId, ConcurrentHashMap())
     }
 
@@ -70,7 +70,7 @@ internal abstract class MessageQueue<V> {
     protected abstract fun isCloseCommand(message: V): Boolean
 
     private fun poll(localId: Int, command: Int): V? {
-        val streamQueues = queues[localId] ?: throw IllegalStateException("Not listening for localId: $localId")
+        val streamQueues = queues[localId] ?: throw AdbStreamClosed(localId)
         val message = streamQueues[command]?.poll()
         if (message == null && !openStreams.contains(localId)) {
             throw AdbStreamClosed(localId)
diff --git a/dadb/src/test/kotlin/dadb/MessageQueueTest.kt b/dadb/src/test/kotlin/dadb/MessageQueueTest.kt
index 30e914e..0cf0072 100644
--- a/dadb/src/test/kotlin/dadb/MessageQueueTest.kt
+++ b/dadb/src/test/kotlin/dadb/MessageQueueTest.kt
@@ -22,6 +22,7 @@ import java.util.concurrent.LinkedBlockingDeque
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.Test
 import kotlin.test.BeforeTest
+import kotlin.test.assertFailsWith
 
 internal class MessageQueueTest : BaseConcurrencyTest() {
 
@@ -48,6 +49,28 @@ internal class MessageQueueTest : BaseConcurrencyTest() {
         take(1, 0)
     }
 
+    @Test
+    fun startListening_rejectsLocalIdAlreadyInUse() {
+        // Two streams listening on the same localId share a single queues entry.
+        // The first one to close removes that entry and the survivor's next
+        // take() dies with "Not listening for localId", even though it was
+        // never closed. Registering a duplicate must fail loudly instead.
+        assertFailsWith<IllegalStateException> {
+            messageQueue.startListening(0)
+        }
+    }
+
+    @Test
+    fun take_afterStopListening_throwsAdbStreamClosed() {
+        messageQueue.stopListening(1)
+
+        // A stream that is gone should surface as the graceful AdbStreamClosed
+        // (which callers already handle), not IllegalStateException.
+        assertFailsWith<AdbStreamClosed> {
+            messageQueue.take(1, 0)
+        }
+    }
+
     @Test
     fun concurrency() {
         val sendsRemaining = AtomicInteger(1000)