From 3fb45158939557bda5e508af07f814c39d02bb88 Mon Sep 17 00:00:00 2001
From: wh201906 <wh201906@yandex.com>
Date: Tue, 26 Dec 2023 20:21:35 +0800
Subject: [PATCH] Check TCP connection state

---
 CHANGELOG.md                 |  1 +
 client/src/uart/uart_posix.c | 17 ++++++++++++++++-
 client/src/uart/uart_win32.c | 17 ++++++++++++++++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1b4d12182..3877e21f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@ All notable changes to this project will be documented in this file.
 This project uses the changelog in accordance with [keepchangelog](http://keepachangelog.com/). Please use this to write notable changes, which is not the same as git commit log...
 
 ## [unreleased][unreleased]
+ - Changed `uart_receive()` - Check if TCP connection is lost (@wh201906)
  - Change `data detectclock` - now tries all clocks if called w/o any params (@iceman1001)
  - Changed `lf search -1u` - improved the autocorrelation detection for unknown signals (@iceman1001)
  - Fixed `hf emrtd dump` stack smashing on device side (@iceman1001)
diff --git a/client/src/uart/uart_posix.c b/client/src/uart/uart_posix.c
index 588e9ae90..2a74df483 100644
--- a/client/src/uart/uart_posix.c
+++ b/client/src/uart/uart_posix.c
@@ -69,6 +69,7 @@ struct timeval timeout = {
 
 static uint32_t newtimeout_value = 0;
 static bool newtimeout_pending = false;
+static uint8_t rx_empty_counter = 0;
 
 int uart_reconfigure_timeouts(uint32_t value) {
     newtimeout_value = value;
@@ -727,7 +728,21 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
         // Retrieve the count of the incoming bytes
         res = ioctl(spu->fd, FIONREAD, &byteCount);
 //        PrintAndLogEx(ERR, "UART:: RX ioctl res %d byteCount %u", res, byteCount);
-        if (res < 0) return PM3_ENOTTY;
+        if (res < 0) {
+            // error occurred (maybe disconnected)
+            // This happens when USB-CDC connection is lost
+            return PM3_ENOTTY;
+        } else if (byteCount == 0) {
+            // select() > 0 && byteCount > 0 ===> data available
+            // select() > 0 && byteCount always equals to 0 ===> maybe disconnected
+            // This happens when TCP connection is lost
+            rx_empty_counter++;
+            if (rx_empty_counter > 3) {
+                return PM3_ENOTTY;
+            }
+        } else {
+            rx_empty_counter = 0;
+        }
 
         // For UDP connection, put the incoming data into the buffer and handle them in the next round
         if (spu->udpBuffer != NULL) {
diff --git a/client/src/uart/uart_win32.c b/client/src/uart/uart_win32.c
index 43299ca5a..b2e6a0adc 100644
--- a/client/src/uart/uart_win32.c
+++ b/client/src/uart/uart_win32.c
@@ -49,6 +49,7 @@ struct timeval timeout = {
 
 uint32_t newtimeout_value = 0;
 bool newtimeout_pending = false;
+uint8_t rx_empty_counter = 0;
 
 int uart_reconfigure_timeouts(uint32_t value) {
     newtimeout_value = value;
@@ -679,7 +680,21 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
             // Retrieve the count of the incoming bytes
             res = ioctlsocket(spw->hSocket, FIONREAD, (u_long *)&byteCount);
             // PrintAndLogEx(ERR, "UART:: RX ioctl res %d byteCount %u", res, byteCount);
-            if (res == SOCKET_ERROR) return PM3_ENOTTY;
+            if (res == SOCKET_ERROR) {
+                // error occurred (maybe disconnected)
+                // This happens when USB-CDC connection is lost
+                return PM3_ENOTTY;
+            } else if (byteCount == 0) {
+                // select() > 0 && byteCount > 0 ===> data available
+                // select() > 0 && byteCount always equals to 0 ===> maybe disconnected
+                // This happens when TCP connection is lost
+                rx_empty_counter++;
+                if (rx_empty_counter > 3) {
+                    return PM3_ENOTTY;
+                }
+            } else {
+                rx_empty_counter = 0;
+            }
 
             // For UDP connection, put the incoming data into the buffer and handle them in the next round
             if (spw->udpBuffer != NULL) {