unify connection tracker for ble and wifi (#537)

* unify connection tracker for ble and wifi * cleanup nicely
2026-07-24 23:25:19 +00:00 · 2026-01-04 02:00:43 +07:00
parent c12af45b6b
commit 12eb3e5a91
4 changed files with 297 additions and 169 deletions
@@ -16,14 +16,11 @@ import java.util.concurrent.CopyOnWriteArrayList
 class BluetoothConnectionTracker(
    private val connectionScope: CoroutineScope,
    private val powerManager: PowerManager
-) {
+) : MeshConnectionTracker(connectionScope, TAG) {
    
    companion object {
        private const val TAG = "BluetoothConnectionTracker"
-        private const val CONNECTION_RETRY_DELAY = com.bitchat.android.util.AppConstants.Mesh.CONNECTION_RETRY_DELAY_MS
-        private const val MAX_CONNECTION_ATTEMPTS = com.bitchat.android.util.AppConstants.Mesh.MAX_CONNECTION_ATTEMPTS
        private const val CLEANUP_DELAY = com.bitchat.android.util.AppConstants.Mesh.CONNECTION_CLEANUP_DELAY_MS
-        private const val CLEANUP_INTERVAL = com.bitchat.android.util.AppConstants.Mesh.CONNECTION_CLEANUP_INTERVAL_MS // 30 seconds
    }
    
    // Connection tracking - reduced memory footprint
@@ -36,12 +33,6 @@ class BluetoothConnectionTracker(
    // RSSI tracking from scan results (for devices we discover but may connect as servers)
    private val scanRSSI = ConcurrentHashMap<String, Int>()
    
-    // Connection attempt tracking with automatic cleanup
-    private val pendingConnections = ConcurrentHashMap<String, ConnectionAttempt>()
-    
-    // State management
-    private var isActive = false
-    
    /**
     * Consolidated device connection information
     */
@@ -54,37 +45,28 @@ class BluetoothConnectionTracker(
        val connectedAt: Long = System.currentTimeMillis()
    )
    
-    /**
-     * Connection attempt tracking with automatic expiry
-     */
-    data class ConnectionAttempt(
-        val attempts: Int,
-        val lastAttempt: Long = System.currentTimeMillis()
-    ) {
-        fun isExpired(): Boolean = 
-            System.currentTimeMillis() - lastAttempt > CONNECTION_RETRY_DELAY * 2
-        
-        fun shouldRetry(): Boolean = 
-            attempts < MAX_CONNECTION_ATTEMPTS && 
-            System.currentTimeMillis() - lastAttempt > CONNECTION_RETRY_DELAY
+    override fun start() {
+        super.start()
    }
    
-    /**
-     * Start the connection tracker
-     */
-    fun start() {
-        isActive = true
-        startPeriodicCleanup()
-    }
-    
-    /**
-     * Stop the connection tracker
-     */
-    fun stop() {
-        isActive = false
+    override fun stop() {
+        super.stop()
        cleanupAllConnections()
        clearAllConnections()
    }
+
+    // Abstract implementations
+    override fun isConnected(id: String): Boolean = connectedDevices.containsKey(id)
+    
+    override fun disconnect(id: String) {
+        connectedDevices[id]?.gatt?.let {
+            try { it.disconnect() } catch (_: Exception) { }
+        }
+        cleanupDeviceConnection(id)
+        Log.d(TAG, "Requested disconnect for $id")
+    }
+
+    override fun getConnectionCount(): Int = connectedDevices.size
    
    /**
     * Add a device connection
@@ -92,7 +74,7 @@ class BluetoothConnectionTracker(
    fun addDeviceConnection(deviceAddress: String, deviceConn: DeviceConnection) {
        Log.d(TAG, "Tracker: Adding device connection for $deviceAddress (isClient: ${deviceConn.isClient}")
        connectedDevices[deviceAddress] = deviceConn
-        pendingConnections.remove(deviceAddress)
+        removePendingConnection(deviceAddress)
        // Mark as awaiting first ANNOUNCE on this connection
        firstAnnounceSeen[deviceAddress] = false
    }
@@ -167,67 +149,17 @@ class BluetoothConnectionTracker(
    /**
     * Check if device is already connected
     */
-    fun isDeviceConnected(deviceAddress: String): Boolean {
-        return connectedDevices.containsKey(deviceAddress)
-    }
-    
-    /**
-     * Check if connection attempt is allowed
-     */
-    fun isConnectionAttemptAllowed(deviceAddress: String): Boolean {
-        val existingAttempt = pendingConnections[deviceAddress]
-        return existingAttempt?.let { 
-            it.isExpired() || it.shouldRetry() 
-        } ?: true
-    }
-    
-    /**
-     * Add a pending connection attempt
-     */
-    fun addPendingConnection(deviceAddress: String): Boolean {
-        Log.d(TAG, "Tracker: Adding pending connection for $deviceAddress")
-        synchronized(pendingConnections) {
-            // Double-check inside synchronized block
-            val currentAttempt = pendingConnections[deviceAddress]
-            if (currentAttempt != null && !currentAttempt.isExpired() && !currentAttempt.shouldRetry()) {
-                Log.d(TAG, "Tracker: Connection attempt already in progress for $deviceAddress")
-                return false
-            }
-            if (currentAttempt != null) {
-                Log.d(TAG, "Tracker: current attempt: $currentAttempt")
-            }
-            
-            // Update connection attempt atomically
-            // If the previous attempt window expired, reset backoff to 1; otherwise increment
-            val attempts = if (currentAttempt?.isExpired() == true) 1 else (currentAttempt?.attempts ?: 0) + 1
-            pendingConnections[deviceAddress] = ConnectionAttempt(attempts)
-            Log.d(TAG, "Tracker: Added pending connection for $deviceAddress (attempts: $attempts)")
-            return true
-        }
-    }
+    fun isDeviceConnected(deviceAddress: String): Boolean = isConnected(deviceAddress)
    
    /**
     * Disconnect a specific device (by MAC address)
     */
-    fun disconnectDevice(deviceAddress: String) {
-        connectedDevices[deviceAddress]?.gatt?.let {
-            try { it.disconnect() } catch (_: Exception) { }
-        }
-        cleanupDeviceConnection(deviceAddress)
-        Log.d(TAG, "Requested disconnect for $deviceAddress")
-    }
-
-    /**
-     * Remove a pending connection
-     */
-    fun removePendingConnection(deviceAddress: String) {
-        pendingConnections.remove(deviceAddress)
-    }
+    fun disconnectDevice(deviceAddress: String) = disconnect(deviceAddress)
    
    /**
     * Get connected device count
     */
-    fun getConnectedDeviceCount(): Int = connectedDevices.size
+    fun getConnectedDeviceCount(): Int = getConnectionCount()
    
    /**
     * Check if connection limit is reached
@@ -335,36 +267,6 @@ class BluetoothConnectionTracker(
        return firstAnnounceSeen[deviceAddress] == true
    }
    
-    /**
-     * Start periodic cleanup of expired connections
-     */
-    private fun startPeriodicCleanup() {
-        connectionScope.launch {
-            while (isActive) {
-                delay(CLEANUP_INTERVAL)
-                
-                if (!isActive) break
-                
-                try {
-                    // Clean up expired pending connections
-                    val expiredConnections = pendingConnections.filter { it.value.isExpired() }
-                    expiredConnections.keys.forEach { pendingConnections.remove(it) }
-                    
-                    // Log cleanup if any
-                    if (expiredConnections.isNotEmpty()) {
-                        Log.d(TAG, "Cleaned up ${expiredConnections.size} expired connection attempts")
-                    }
-                    
-                    // Log current state
-                    Log.d(TAG, "Periodic cleanup: ${connectedDevices.size} connections, ${pendingConnections.size} pending")
-                    
-                } catch (e: Exception) {
-                    Log.w(TAG, "Error in periodic cleanup: ${e.message}")
-                }
-            }
-        }
-    }
-    
    /**
     * Get debug information
     */
@@ -0,0 +1,140 @@
+package com.bitchat.android.mesh
+
+import android.util.Log
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * Abstract base tracker for mesh connections (BLE, Wi-Fi Aware, etc.)
+ * Encapsulates common state machine logic:
+ * - Connection attempt tracking (retries, backoff)
+ * - Pending connection management
+ * - Automatic cleanup of expired attempts
+ */
+abstract class MeshConnectionTracker(
+    private val scope: CoroutineScope,
+    protected val tag: String
+) {
+    companion object {
+        const val CONNECTION_RETRY_DELAY = 5_000L
+        const val MAX_CONNECTION_ATTEMPTS = 3
+        const val CLEANUP_INTERVAL = 30_000L
+    }
+
+    /**
+     * Connection attempt tracking with automatic expiry
+     */
+    protected data class ConnectionAttempt(
+        val attempts: Int,
+        val lastAttempt: Long = System.currentTimeMillis()
+    ) {
+        fun isExpired(): Boolean =
+            System.currentTimeMillis() - lastAttempt > CONNECTION_RETRY_DELAY * 2
+
+        fun shouldRetry(): Boolean =
+            attempts < MAX_CONNECTION_ATTEMPTS &&
+                    System.currentTimeMillis() - lastAttempt > CONNECTION_RETRY_DELAY
+    }
+
+    // Tracks in-progress or failed attempts
+    protected val pendingConnections = ConcurrentHashMap<String, ConnectionAttempt>()
+
+    private var isActive = false
+
+    /**
+     * Start the tracker and its cleanup loop
+     */
+    open fun start() {
+        isActive = true
+        startPeriodicCleanup()
+    }
+
+    /**
+     * Stop the tracker
+     */
+    open fun stop() {
+        isActive = false
+        pendingConnections.clear()
+    }
+
+    /**
+     * Check if a connection attempt is allowed for this peer/address
+     */
+    fun isConnectionAttemptAllowed(id: String): Boolean {
+        // If already connected, usually no need to retry (subclasses can override logic if needed,
+        // but typically the caller checks isConnected() first).
+        
+        val existingAttempt = pendingConnections[id]
+        return existingAttempt?.let {
+            it.isExpired() || it.shouldRetry()
+        } ?: true
+    }
+
+    /**
+     * Record a new connection attempt.
+     * Returns true if the attempt was recorded (allowed), false if skipped.
+     */
+    fun addPendingConnection(id: String): Boolean {
+        synchronized(pendingConnections) {
+            val currentAttempt = pendingConnections[id]
+            
+            // If strictly not allowed right now, reject
+            if (currentAttempt != null && !currentAttempt.isExpired() && !currentAttempt.shouldRetry()) {
+                Log.d(tag, "Connection attempt already in progress for $id")
+                return false
+            }
+
+            // Update attempt count
+            // Reset to 1 if expired, otherwise increment
+            val attempts = if (currentAttempt?.isExpired() == true) 1 else (currentAttempt?.attempts ?: 0) + 1
+            pendingConnections[id] = ConnectionAttempt(attempts)
+            Log.d(tag, "Added pending connection for $id (attempts: $attempts)")
+            return true
+        }
+    }
+
+    /**
+     * Remove a pending attempt (e.g., on success or fatal error)
+     */
+    fun removePendingConnection(id: String) {
+        pendingConnections.remove(id)
+    }
+
+    /**
+     * Abstract: Subclasses must define what "connected" means
+     */
+    abstract fun isConnected(id: String): Boolean
+
+    /**
+     * Abstract: Subclasses must implement disconnect logic
+     */
+    abstract fun disconnect(id: String)
+
+    /**
+     * Abstract: Subclasses report their active connection count
+     */
+    abstract fun getConnectionCount(): Int
+
+    private fun startPeriodicCleanup() {
+        scope.launch {
+            while (isActive) {
+                try {
+                    delay(CLEANUP_INTERVAL)
+                    if (!isActive) break
+
+                    // Clean up expired pending connections
+                    val expired = pendingConnections.filter { it.value.isExpired() }
+                    expired.keys.forEach { pendingConnections.remove(it) }
+
+                    if (expired.isNotEmpty()) {
+                        Log.d(tag, "Cleaned up ${expired.size} expired connection attempts")
+                    }
+                } catch (e: Exception) {
+                    Log.w(tag, "Error in periodic cleanup: ${e.message}")
+                }
+            }
+        }
+    }
+}
@@ -0,0 +1,95 @@
+package com.bitchat.android.wifiaware
+
+import android.net.ConnectivityManager
+import android.util.Log
+import com.bitchat.android.mesh.MeshConnectionTracker
+import kotlinx.coroutines.CoroutineScope
+import java.net.ServerSocket
+import java.net.Socket
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * Tracks Wi-Fi Aware connections and manages retry logic using the shared state machine.
+ */
+class WifiAwareConnectionTracker(
+    scope: CoroutineScope,
+    private val cm: ConnectivityManager
+) : MeshConnectionTracker(scope, TAG) {
+
+    companion object {
+        private const val TAG = "WifiAwareConnectionTracker"
+    }
+
+    // Active resources per peer
+    val peerSockets = ConcurrentHashMap<String, Socket>()
+    val serverSockets = ConcurrentHashMap<String, ServerSocket>()
+    val networkCallbacks = ConcurrentHashMap<String, ConnectivityManager.NetworkCallback>()
+
+    override fun isConnected(id: String): Boolean {
+        // We consider it connected if we have a client socket to them
+        return peerSockets.containsKey(id)
+    }
+
+    override fun disconnect(id: String) {
+        Log.d(TAG, "Disconnecting peer $id")
+        
+        // 1. Close client socket
+        peerSockets.remove(id)?.let {
+            try { it.close() } catch (e: Exception) { Log.w(TAG, "Error closing socket for $id: ${e.message}") }
+        }
+
+        // 2. Close server socket
+        serverSockets.remove(id)?.let {
+            try { it.close() } catch (e: Exception) { Log.w(TAG, "Error closing server socket for $id: ${e.message}") }
+        }
+
+        // 3. Unregister network callback
+        networkCallbacks.remove(id)?.let {
+            try { cm.unregisterNetworkCallback(it) } catch (e: Exception) { Log.w(TAG, "Error unregistering callback for $id: ${e.message}") }
+        }
+    }
+
+    override fun getConnectionCount(): Int = peerSockets.size
+
+    /**
+     * Successfully established a client connection
+     */
+    fun onClientConnected(peerId: String, socket: Socket) {
+        peerSockets[peerId] = socket
+        removePendingConnection(peerId) // Clear retry state on success
+    }
+
+    fun addServerSocket(peerId: String, socket: ServerSocket) {
+        serverSockets[peerId] = socket
+    }
+
+    fun addNetworkCallback(peerId: String, callback: ConnectivityManager.NetworkCallback) {
+        networkCallbacks[peerId] = callback
+    }
+    
+    /**
+     * Clean up all resources
+     */
+    override fun stop() {
+        super.stop()
+        val allIds = peerSockets.keys + serverSockets.keys + networkCallbacks.keys
+        allIds.toSet().forEach { disconnect(it) }
+    }
+    
+    fun getDebugInfo(): String {
+        return buildString {
+            appendLine("Aware Connections: ${getConnectionCount()}")
+            peerSockets.keys.forEach { pid ->
+                appendLine("  - $pid (Socket)")
+            }
+            appendLine("Server Sockets: ${serverSockets.size}")
+            serverSockets.keys.forEach { pid ->
+                appendLine("  - $pid (Listening)")
+            }
+            appendLine("Pending Attempts: ${pendingConnections.size}")
+            pendingConnections.forEach { (pid, attempt) ->
+                appendLine("  - $pid: ${attempt.attempts} attempts")
+            }
+        }
+    }
+}
@@ -89,21 +89,18 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
    // Delegate
    var delegate: WifiAwareMeshDelegate? = null

-    // Transport state
-    private val peerSockets = ConcurrentHashMap<String, Socket>()
-    private val serverSockets = ConcurrentHashMap<String, ServerSocket>()
-    private val networkCallbacks = ConcurrentHashMap<String, ConnectivityManager.NetworkCallback>()
+    // Coroutines - must be initialized before tracker
+    private val serviceScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
    private val cm = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
+
+    // Transport state
+    private val connectionTracker = WifiAwareConnectionTracker(serviceScope, cm)
    private val handleToPeerId = ConcurrentHashMap<PeerHandle, String>() // discovery mapping
    private val discoveredTimestamps = ConcurrentHashMap<String, Long>() // peerID -> last seen time
-    private val connectionAttempts = ConcurrentHashMap<String, Long>() // peerID -> last connection attempt time

    // Timestamp dedupe
    private val lastTimestamps = ConcurrentHashMap<String, ULong>()

-    // Coroutines
-    private val serviceScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
-
    init {
        setupDelegates()
        messageHandler.packetProcessor = packetProcessor
@@ -165,7 +162,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
     */
    private fun broadcastRaw(bytes: ByteArray) {
        var sent = 0
-        peerSockets.forEach { (pid, sock) ->
+        connectionTracker.peerSockets.forEach { (pid, sock) ->
            try {
                sock.getOutputStream().write(bytes)
                sent++
@@ -215,7 +212,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        // Wi-Fi Aware uses full packets; no fragmentation
        val data = packet.toBinaryData() ?: return
        serviceScope.launch {
-            val sock = peerSockets[peerID]
+            val sock = connectionTracker.peerSockets[peerID]
            if (sock == null) {
                Log.w(TAG, "TX: no socket for ${peerID.take(8)}")
                return@launch
@@ -532,6 +529,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic

        sendPeriodicBroadcastAnnounce()
        startPeriodicConnectionMaintenance()
+        connectionTracker.start()
        gossipSyncManager.start()
        
        // Register with cross-layer transport bridge
@@ -555,18 +553,13 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
            delay(200)

            gossipSyncManager.stop()
+            connectionTracker.stop() // Handles socket closing and callback unregistration

-            networkCallbacks.values.forEach { runCatching { cm.unregisterNetworkCallback(it) } }
-            networkCallbacks.clear()
            publishSession?.close();   publishSession   = null
            subscribeSession?.close(); subscribeSession = null
            wifiAwareSession?.close(); wifiAwareSession = null

-            serverSockets.values.forEach { it.closeQuietly() }
-            peerSockets.values.forEach { it.closeQuietly() }
            handleToPeerId.clear()
-            serverSockets.clear()
-            peerSockets.clear()

            peerManager.shutdown()
            fragmentManager.shutdown()
@@ -608,9 +601,9 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                        (now - ts) < 5 * 60 * 1000 // Seen in last 5 minutes
                    }.keys

-                    // 2. Filter out those who are already connected (either as client or server)
+                    // 2. Filter out those who are already connected
                    val disconnectedPeers = recentDiscovered.filter { peerId ->
-                        !peerSockets.containsKey(peerId) && !serverSockets.containsKey(peerId)
+                        !connectionTracker.isConnected(peerId)
                    }

                    // 3. Attempt reconnection
@@ -618,19 +611,18 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                        // Find the PeerHandle for this peerId
                        val handle = handleToPeerId.entries.find { it.value == peerId }?.key ?: continue
                        
-                        // Rate limit attempts: max 1 per minute per peer
-                        val lastAttempt = connectionAttempts[peerId] ?: 0L
-                        if (now - lastAttempt < 60_000) continue
+                        // Check tracker policy
+                        if (!connectionTracker.isConnectionAttemptAllowed(peerId)) continue

                        Log.i(TAG, "🔄 Maintenance: attempting reconnect to ${peerId.take(8)}")
-                        connectionAttempts[peerId] = now
-                        
-                        // Resend ping to trigger handshake
-                        val msgId = (System.nanoTime() and 0x7fffffff).toInt()
-                        try {
-                            subscribeSession?.sendMessage(handle, msgId, myPeerID.toByteArray())
-                        } catch (e: Exception) {
-                            Log.w(TAG, "Failed to send maintenance ping to ${peerId.take(8)}: ${e.message}")
+                        if (connectionTracker.addPendingConnection(peerId)) {
+                            // Resend ping to trigger handshake
+                            val msgId = (System.nanoTime() and 0x7fffffff).toInt()
+                            try {
+                                subscribeSession?.sendMessage(handle, msgId, myPeerID.toByteArray())
+                            } catch (e: Exception) {
+                                Log.w(TAG, "Failed to send maintenance ping to ${peerId.take(8)}: ${e.message}")
+                            }
                        }
                    }
                } catch (e: Exception) {
@@ -654,13 +646,13 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        val peerId = handleToPeerId[peerHandle] ?: return
        if (!amIServerFor(peerId)) return

-        if (serverSockets.containsKey(peerId)) {
+        if (connectionTracker.serverSockets.containsKey(peerId)) {
            Log.v(TAG, "↪ already serving $peerId, skipping")
            return
        }

        val ss = ServerSocket(0)
-        serverSockets[peerId] = ss
+        connectionTracker.addServerSocket(peerId, ss)
        val port = ss.localPort
        
        // Ensure port is set to reuse if connection was recently closed (TIME_WAIT)
@@ -688,7 +680,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                    try { network.bindSocket(client) } catch (e: Exception) { Log.w(TAG, "Server bindSocket EPERM: ${e.message}") }
                    client.keepAlive = true
                    Log.d(TAG, "SERVER: accepted TCP from ${peerId.take(8)} addr=${client.inetAddress?.hostAddress}")
-                    peerSockets[peerId] = client
+                    connectionTracker.onClientConnected(peerId, client)
                    try { peerManager.setDirectConnection(peerId, true) } catch (_: Exception) {}
                    try { peerManager.addOrUpdatePeer(peerId, peerId) } catch (_: Exception) {}
                    listenerExec.execute { listenToPeer(client, peerId) }
@@ -705,12 +697,12 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                }
            }
            override fun onLost(network: Network) {
-                networkCallbacks.remove(peerId)
+                connectionTracker.networkCallbacks.remove(peerId)
                Log.d(TAG, "SERVER: network lost for ${peerId.take(8)}")
            }
        }

-        networkCallbacks[peerId] = cb
+        connectionTracker.addNetworkCallback(peerId, cb)
        Log.d(TAG, "SERVER: requesting Aware network for ${peerId.take(8)}")
        cm.requestNetwork(req, cb)

@@ -745,7 +737,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        serviceScope.launch {
            try {
                val os = client.getOutputStream()
-                while (peerSockets.containsKey(peerId)) {
+                while (connectionTracker.isConnected(peerId)) {
                    try { os.write(0) } catch (_: IOException) { break }
                    delay(2_000)
                }
@@ -754,7 +746,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        // Discovery keep-alive
        serviceScope.launch {
            var msgId = 0
-            while (peerSockets.containsKey(peerId)) {
+            while (connectionTracker.isConnected(peerId)) {
                try { pubSession.sendMessage(peerHandle, msgId++, ByteArray(0)) } catch (_: Exception) { break }
                delay(20_000)
            }
@@ -776,7 +768,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic

        val peerId = handleToPeerId[peerHandle] ?: return
        if (amIServerFor(peerId)) return
-        if (peerSockets.containsKey(peerId)) {
+        if (connectionTracker.peerSockets.containsKey(peerId)) {
            Log.v(TAG, "↪ already client-connected to $peerId, skipping")
            return
        }
@@ -797,7 +789,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                // Do not bind process for Aware; use per-socket binding instead
            }
            override fun onCapabilitiesChanged(network: Network, nc: NetworkCapabilities) {
-                if (peerSockets.containsKey(peerId)) return
+                if (connectionTracker.peerSockets.containsKey(peerId)) return
                val info = (nc.transportInfo as? WifiAwareNetworkInfo) ?: return
                val addr = info.peerIpv6Addr as? Inet6Address ?: return

@@ -824,7 +816,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                    sock.connect(java.net.InetSocketAddress(scopedAddr, port), 7000)
                    Log.d(TAG, "CLIENT: TCP connected to ${peerId.take(8)} addr=$scopedAddr:$port (iface=$iface)")

-                    peerSockets[peerId] = sock
+                    connectionTracker.onClientConnected(peerId, sock)
                    try { peerManager.setDirectConnection(peerId, true) } catch (_: Exception) {}
                    try { peerManager.addOrUpdatePeer(peerId, peerId) } catch (_: Exception) {}
                    listenerExec.execute { listenToPeer(sock, peerId) }
@@ -841,12 +833,12 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
                }
            }
            override fun onLost(network: Network) {
-                networkCallbacks.remove(peerId)
+                connectionTracker.networkCallbacks.remove(peerId)
                Log.d(TAG, "CLIENT: network lost for ${peerId.take(8)}")
            }
        }

-        networkCallbacks[peerId] = cb
+        connectionTracker.addNetworkCallback(peerId, cb)
        Log.d(TAG, "CLIENT: requesting Aware network for ${peerId.take(8)}")
        cm.requestNetwork(req, cb)
    }
@@ -863,7 +855,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        serviceScope.launch {
            try {
                val os = sock.getOutputStream()
-                while (peerSockets.containsKey(peerId)) {
+                while (connectionTracker.isConnected(peerId)) {
                    try { os.write(0) } catch (_: IOException) { break }
                    delay(2_000)
                }
@@ -872,7 +864,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        // Discovery keep-alive
        serviceScope.launch {
            var msgId = 0
-            while (peerSockets.containsKey(peerId)) {
+            while (connectionTracker.isConnected(peerId)) {
                try { subscribeSession?.sendMessage(peerHandle, msgId++, ByteArray(0)) } catch (_: Exception) { break }
                delay(20_000)
            }
@@ -929,9 +921,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
        serviceScope.launch {
            Log.d(TAG, "Cleaning up peer: $initialId")
            
-            peerSockets.remove(initialId)?.closeQuietly()
-            serverSockets.remove(initialId)?.closeQuietly()
-            networkCallbacks.remove(initialId)?.let { runCatching { cm.unregisterNetworkCallback(it) } }
+            connectionTracker.disconnect(initialId)
            peerManager.removePeer(initialId)
        }
    }
@@ -1257,7 +1247,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
     * Prefers the scoped IPv6 address format.
     */
    fun getDeviceAddressForPeer(peerID: String): String? =
-        peerSockets[peerID]?.let { resolveScopedAddress(it) }
+        connectionTracker.peerSockets[peerID]?.let { resolveScopedAddress(it) }

    /**
     * Helper to resolve a scoped IPv6 address from a socket for UI display.
@@ -1281,7 +1271,7 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
     */
    fun getDeviceAddressToPeerMapping(): Map<String, String> {
        val map = mutableMapOf<String, String>()
-        peerSockets.forEach { (pid, sock) ->
+        connectionTracker.peerSockets.forEach { (pid, sock) ->
            map[pid] = resolveScopedAddress(sock) ?: "unknown"
        }
        return map
@@ -1308,7 +1298,8 @@ class WifiAwareMeshService(private val context: Context) : TransportBridgeServic
    fun getDebugStatus(): String = buildString {
        appendLine("=== Wi-Fi Aware Mesh Debug Status ===")
        appendLine("My Peer ID: $myPeerID")
-        appendLine("Peers: ${peerSockets.keys}")
+        appendLine("Peers: ${connectionTracker.peerSockets.keys}")
+        appendLine(connectionTracker.getDebugInfo())
        appendLine(peerManager.getDebugInfo(getDeviceAddressToPeerMapping()))
        appendLine(fragmentManager.getDebugInfo())
        appendLine(securityManager.getDebugInfo())