fix: wifi aware mesh private dms (#561)

* fix(wifi-aware): restore peer check and remove aggressive session restarts

* synchronized access to sockets.

* clear old peer socket `onClientConnected` if one was found.

* address second round of review

* fix: Restart publish session on termination in WifiAwareMeshService

* fix: cleanup connection tracker resources when handling peer disconnection without active socket

---------

Co-authored-by: aidenvalu3 <>

app/src/main/java/com/bitchat/android/noise/NoiseSession.kt

@@ -346,7 +346,8 @@ class NoiseSession(
                 Log.w(TAG, "Warning: XX message 1 size ${firstMessage.size} != expected $XX_MESSAGE_1_SIZE")
             }
             
-            Log.d(TAG, "Sending XX handshake message 1 to $peerID (${firstMessage.size} bytes) currentPattern: $currentPattern")
+            val ePrefix = firstMessage.take(4).toByteArray().toHexString()
+            Log.d(TAG, "Sending XX handshake message 1 to $peerID (${firstMessage.size} bytes) e_prefix=$ePrefix currentPattern: $currentPattern")
             return firstMessage
         } catch (e: Exception) {
             state = NoiseSessionState.Failed(e)
@@ -361,7 +362,8 @@ class NoiseSession(
      */
     @Synchronized
     fun processHandshakeMessage(message: ByteArray): ByteArray? {
-        Log.d(TAG, "Processing handshake message from $peerID (${message.size} bytes)")
+        val inputPrefix = message.take(4).toByteArray().toHexString()
+        Log.d(TAG, "Processing handshake message from $peerID (${message.size} bytes) prefix=$inputPrefix")
         
         try {
             // Initialize as responder if receiving first message
@@ -387,7 +389,8 @@ class NoiseSession(
             // Read the incoming message - the Noise library will handle validation
             val payloadLength = handshakeStateLocal.readMessage(message, 0, message.size, payloadBuffer, 0)
             currentPattern++
-            Log.d(TAG, "Read handshake message, payload length: $payloadLength currentPattern: $currentPattern")
+            val readPrefix = message.take(4).toByteArray().toHexString()
+            Log.d(TAG, "Read handshake message, payload length: $payloadLength prefix=$readPrefix currentPattern: $currentPattern")
             
             // Check what action the handshake state wants us to take next
             val action = handshakeStateLocal.getAction()

app/src/main/java/com/bitchat/android/wifi-aware/SyncedSocket.kt

@@ -0,0 +1,86 @@
+package com.bitchat.android.wifiaware
+
+import java.io.DataInputStream
+import java.io.DataOutputStream
+import java.io.IOException
+import java.net.Socket
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
+import android.util.Log
+
+/**
+ * A synchronized wrapper around a raw Socket that implements a framed protocol:
+ * [4 bytes length][N bytes payload]
+ */
+class SyncedSocket(val rawSocket: Socket) {
+    private val TAG = "SyncedSocket"
+    private val writeLock = ReentrantLock()
+    private val readLock = ReentrantLock()
+    
+    private val inputStream: DataInputStream
+    private val outputStream: DataOutputStream
+
+    init {
+        // We wrap streams to create DataInput/Output helpers
+        inputStream = DataInputStream(rawSocket.getInputStream())
+        outputStream = DataOutputStream(rawSocket.getOutputStream())
+    }
+
+    /**
+     * Writes a framed message to the socket.
+     * Thread-safe.
+     */
+    fun write(data: ByteArray) {
+        writeLock.withLock {
+            Log.v(TAG, "Writing frame of size: ${data.size}")
+            outputStream.writeInt(data.size)
+            if (data.isNotEmpty()) {
+                outputStream.write(data)
+            }
+            outputStream.flush()
+        }
+    }
+
+    /**
+     * Reads a framed message from the socket.
+     * Blocks until a full frame is available.
+     * Returns null if socket is closed or EOF.
+     * Returns empty byte array for keep-alive (0 length frame).
+     */
+    fun read(): ByteArray? {
+        readLock.withLock {
+            try {
+                // Read length prefix
+                val length = try {
+                    inputStream.readInt()
+                } catch (e: java.io.EOFException) {
+                    return null
+                }
+                Log.v(TAG, "Reading frame of size: $length")
+                
+                if (length < 0) throw IOException("Negative frame length: $length")
+                if (length > 64 * 1024) throw IOException("Frame length exceeds 64KB limit: $length")
+                
+                if (length == 0) {
+                    return ByteArray(0)
+                }
+
+                val buf = ByteArray(length)
+                inputStream.readFully(buf)
+                return buf
+            } catch (e: IOException) {
+                Log.e(TAG, "Socket read failed: ${e.message}")
+                // Socket closed or error
+                return null
+            }
+        }
+    }
+    
+    fun close() {
+        try { rawSocket.close() } catch (_: Exception) {}
+    }
+
+    fun isClosed() = rawSocket.isClosed
+    fun isConnected() = rawSocket.isConnected
+    val inetAddress get() = rawSocket.inetAddress
+}

app/src/main/java/com/bitchat/android/wifi-aware/WifiAwareConnectionTracker.kt

@@ -21,7 +21,7 @@ class WifiAwareConnectionTracker(
     }
 
     // Active resources per peer
-    val peerSockets = ConcurrentHashMap<String, Socket>()
+    val peerSockets = ConcurrentHashMap<String, SyncedSocket>()
     val serverSockets = ConcurrentHashMap<String, ServerSocket>()
     val networkCallbacks = ConcurrentHashMap<String, ConnectivityManager.NetworkCallback>()
 
@@ -64,7 +64,11 @@ class WifiAwareConnectionTracker(
     /**
      * Successfully established a client connection
      */
-    fun onClientConnected(peerId: String, socket: Socket) {
+    fun onClientConnected(peerId: String, socket: SyncedSocket) {
+        // Close previous socket if one exists to prevent zombie readers
+        peerSockets[peerId]?.let { 
+            try { it.close() } catch (_: Exception) {}
+        }
         peerSockets[peerId] = socket
         removePendingConnection(peerId) // Clear retry state on success
     }

app/src/main/java/com/bitchat/android/wifi-aware/WifiAwareMeshService.kt

@@ -132,7 +132,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
         var sent = 0
         connectionTracker.peerSockets.forEach { (pid, sock) ->
             try {
-                sock.getOutputStream().write(bytes)
+                sock.write(bytes)
                 sent++
             } catch (e: IOException) {
                 Log.e(TAG, "TX: write failed to ${pid.take(8)}: ${e.message}")
@@ -180,7 +180,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
                 return@launch
             }
             try {
-                sock.getOutputStream().write(data)
+                sock.write(data)
                 Log.d(TAG, "TX: packet type=${packet.type} to ${peerID.take(8)} (bytes=${data.size})")
             } catch (e: IOException) {
                 Log.e(TAG, "TX: write to ${peerID.take(8)} failed: ${e.message}")
@@ -461,6 +461,8 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
             .build()
 
         val cb = object : ConnectivityManager.NetworkCallback() {
+            private var activeSocket: SyncedSocket? = null
+
             override fun onAvailable(network: Network) {
                 Log.i(TAG, "SERVER: onAvailable() - Aware network is ready for ${peerId.take(8)}")
                 try {
@@ -469,11 +471,13 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
                     try { network.bindSocket(client) } catch (e: Exception) { Log.w(TAG, "Server bindSocket EPERM: ${e.message}") }
                     client.keepAlive = true
                     Log.i(TAG, "SERVER: Bound and established TCP with ${peerId.take(8)} addr=${client.inetAddress?.hostAddress}")
-                    connectionTracker.onClientConnected(peerId, client)
+                    val synced = SyncedSocket(client)
+                    activeSocket = synced
+                    connectionTracker.onClientConnected(peerId, synced)
                     try { meshCore.setDirectConnection(peerId, true) } catch (_: Exception) {}
                     try { meshCore.addOrUpdatePeer(peerId, peerId) } catch (_: Exception) {}
-                    listenerExec.execute { listenToPeer(client, peerId) }
+                    listenerExec.execute { listenToPeer(synced, peerId) }
-                    handleSubscriberKeepAlive(client, peerId, pubSession, peerHandle)
+                    handleSubscriberKeepAlive(synced, peerId, pubSession, peerHandle)
                     
                     // Kick off Noise handshake for this logical peer
                     if (myPeerID < peerId) {
@@ -493,7 +497,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
             }
 
             override fun onLost(network: Network) {
-                handlePeerDisconnection(peerId)
+                handlePeerDisconnection(peerId, activeSocket)
                 Log.i(TAG, "SERVER: WiFi Aware network lost for ${peerId.take(8)}")
             }
         }
@@ -530,7 +534,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
      * @param peerId ID of the connected peer
      */
     private fun handleSubscriberKeepAlive(
-        client: Socket,
+        client: SyncedSocket,
         peerId: String,
         pubSession: PublishDiscoverySession,
         peerHandle: PeerHandle
@@ -538,9 +542,9 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
         // TCP keep-alive pings
         serviceScope.launch {
             try {
-                val os = client.getOutputStream()
                 while (connectionTracker.isConnected(peerId)) {
-                    try { os.write(0) } catch (_: IOException) { break }
+                    // write empty byte array effectively sends [4 bytes length=0] which is our ping
+                    try { client.write(ByteArray(0)) } catch (_: IOException) { break }
                     delay(2_000)
                 }
             } catch (_: Exception) {}
@@ -587,6 +591,8 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
             .build()
 
         val cb = object : ConnectivityManager.NetworkCallback() {
+            private var activeSocket: SyncedSocket? = null
+
             override fun onAvailable(network: Network) {
                 Log.i(TAG, "CLIENT: onAvailable() - Aware network is ready for ${peerId.take(8)}")
                 // Do not bind process for Aware; use per-socket binding instead
@@ -598,6 +604,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
             }
 
             override fun onCapabilitiesChanged(network: Network, nc: NetworkCapabilities) {
+                if (connectionTracker.peerSockets.containsKey(peerId)) return
                 val info = (nc.transportInfo as? WifiAwareNetworkInfo) ?: return
                 val addr = info.peerIpv6Addr as? Inet6Address ?: return
                 Log.i(TAG, "CLIENT: onCapabilitiesChanged() - Peer IPv6 discovered: $addr")
@@ -625,11 +632,13 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
                     sock.connect(java.net.InetSocketAddress(scopedAddr, port), 7000)
                     Log.i(TAG, "CLIENT: TCP connected to ${peerId.take(8)} at $scopedAddr:$port")
 
-                    connectionTracker.onClientConnected(peerId, sock)
+                    val synced = SyncedSocket(sock)
+                    activeSocket = synced
+                    connectionTracker.onClientConnected(peerId, synced)
                     try { meshCore.setDirectConnection(peerId, true) } catch (_: Exception) {}
                     try { meshCore.addOrUpdatePeer(peerId, peerId) } catch (_: Exception) {}
-                    listenerExec.execute { listenToPeer(sock, peerId) }
+                    listenerExec.execute { listenToPeer(synced, peerId) }
-                    handleServerKeepAlive(sock, peerId, peerHandle)
+                    handleServerKeepAlive(synced, peerId, peerHandle)
                     
                     // Kick off Noise handshake for this logical peer
                     if (myPeerID < peerId) {
@@ -643,7 +652,7 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
                 }
             }
             override fun onLost(network: Network) {
-                handlePeerDisconnection(peerId)
+                handlePeerDisconnection(peerId, activeSocket)
                 Log.i(TAG, "CLIENT: WiFi Aware network lost for ${peerId.take(8)}")
             }
         }
@@ -662,16 +671,15 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
      * Sends periodic TCP and discovery keep-alive messages for server connections.
      */
     private fun handleServerKeepAlive(
-        sock: Socket,
+        sock: SyncedSocket,
         peerId: String,
         peerHandle: PeerHandle
     ) {
         // TCP keep-alive
         serviceScope.launch {
             try {
-                val os = sock.getOutputStream()
                 while (connectionTracker.isConnected(peerId)) {
-                    try { os.write(0) } catch (_: IOException) { break }
+                    try { sock.write(ByteArray(0)) } catch (_: IOException) { break }
                     delay(2_000)
                 }
             } catch (_: Exception) {}
@@ -698,15 +706,15 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
      * @param socket Socket connected to the peer
      * @param initialLogicalPeerId Temporary identifier before peer ID resolution
      */
-    private fun listenToPeer(socket: Socket, initialLogicalPeerId: String) {
+    private fun listenToPeer(socket: SyncedSocket, initialLogicalPeerId: String) {
-        val inStream = socket.getInputStream()
-        val buf = ByteArray(64 * 1024)
-
         while (isActive) {
-            val len = try { inStream.read(buf) } catch (_: IOException) { break }
+            val raw = socket.read() ?: break
-            if (len <= 0) break
+            
+            if (raw.isEmpty()) {
+                // Keep-alive (0 length frame)
+                continue
+            }
 
-            val raw = buf.copyOf(len)
             val pkt = BitchatPacket.fromBinaryData(raw) ?: continue
 
             val senderPeerHex = pkt.senderID?.toHexString()?.take(16) ?: continue
@@ -728,8 +736,8 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
         
         // Breaking out of the loop means the socket is dead or service is stopping.
         Log.i(TAG, "Socket loop terminated for ${initialLogicalPeerId.take(8)} removing peer.")
-        handlePeerDisconnection(initialLogicalPeerId)
+        handlePeerDisconnection(initialLogicalPeerId, socket)
-        socket.closeQuietly()
+        socket.close()
     }
 
     private fun handleNetworkFailure(peerId: String) {
@@ -737,16 +745,32 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
             Log.d(TAG, "Network failure cleanup for: $peerId")
             // Specifically release the callback if it didn't happen automatically
             connectionTracker.releaseNetworkRequest(peerId)
+            
+            if (!connectionTracker.isConnected(peerId)) {
                 meshCore.removePeer(peerId)
+            } else {
+                Log.d(TAG, "Network failure ignored for $peerId - another socket is active")
+            }
         }
     }
 
-    private fun handlePeerDisconnection(initialId: String) {
+    private fun handlePeerDisconnection(initialId: String, socket: SyncedSocket? = null) {
         serviceScope.launch {
-            Log.d(TAG, "Cleaning up peer: $initialId")
+            // Check if this socket is the current active one before nuking the session
-            
+            val currentSocket = connectionTracker.peerSockets[initialId]
+            if (currentSocket === socket) {
+                Log.d(TAG, "Cleaning up peer: $initialId (active socket)")
                 connectionTracker.disconnect(initialId)
                 meshCore.removePeer(initialId)
+            } else if (socket == null && currentSocket == null) {
+                // Fallback: If we don't have a specific socket context but we are already disconnected, ensure cleanup
+                Log.d(TAG, "Cleaning up peer: $initialId (no active socket)")
+                connectionTracker.disconnect(initialId)
+                meshCore.removePeer(initialId)
+            } else {
+                Log.d(TAG, "Ignored disconnection for $initialId - socket replaced or inactive")
+                // Do not remove peer/session, as a new socket has likely taken over
+            }
         }
     }
 
@@ -961,35 +985,30 @@ class WifiAwareMeshService(private val context: Context) : MeshService, Transpor
         meshCore.clearAllEncryptionData()
     }
 
-    /** Utility extension to safely close sockets. */
-    private fun Socket.closeQuietly() = try { close() } catch (_: Exception) {}
-
     /** Utility extension to safely close server sockets. */
     private fun ServerSocket.closeQuietly() = try { close() } catch (_: Exception) {}
 
+
     private inner class WifiAwareTransport : MeshTransport {
         override val id: String = "WIFI"
 
         override fun broadcastPacket(routed: RoutedPacket) {
             this@WifiAwareMeshService.broadcastPacket(routed)
         }
-
         override fun sendPacketToPeer(peerID: String, packet: BitchatPacket) {
             this@WifiAwareMeshService.sendPacketToPeer(peerID, packet)
         }
-
         override fun getDeviceAddressForPeer(peerID: String): String? {
-            return connectionTracker.peerSockets[peerID]?.let { resolveScopedAddress(it) }
+            return connectionTracker.peerSockets[peerID]?.let { resolveScopedAddress(it.rawSocket) }
         }
 
         override fun getDeviceAddressToPeerMapping(): Map<String, String> {
             val map = mutableMapOf<String, String>()
             connectionTracker.peerSockets.forEach { (pid, sock) ->
-                map[pid] = resolveScopedAddress(sock) ?: "unknown"
+                map[pid] = resolveScopedAddress(sock.rawSocket) ?: "unknown"
             }
             return map
         }
-
         override fun getTransportDebugInfo(): String {
             return connectionTracker.getDebugInfo()
         }