From b02d9d4d635056623b7cf132c5412e0574c6382d Mon Sep 17 00:00:00 2001
From: sepfy <sepfy95@gmail.com>
Date: Sun, 15 Sep 2024 22:55:44 +0800
Subject: [PATCH] Use recv timeout func in mbedtls bio to fix EOL issue

---
 src/socket.c        | 45 +++++++--------------------------------------
 src/ssl_transport.c | 24 +++++++++++++++++++++---
 2 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/src/socket.c b/src/socket.c
index ef705a7..2d5efaa 100644
--- a/src/socket.c
+++ b/src/socket.c
@@ -253,40 +253,22 @@ void tcp_socket_close(TcpSocket* tcp_socket) {
 }
 
 int tcp_socket_send(TcpSocket* tcp_socket, const uint8_t* buf, int len) {
-  fd_set write_set;
-  struct timeval tv;
-  int ret = -1;
+  int ret;
 
   if (tcp_socket->fd < 0) {
     LOGE("sendto before socket init");
     return -1;
   }
 
-  FD_ZERO(&write_set);
-  FD_SET(tcp_socket->fd, &write_set);
-
-  tv.tv_sec = 0;
-  tv.tv_usec = 500000;
-
-  if ((ret = select(tcp_socket->fd + 1, NULL, &write_set, NULL, &tv)) < 0) {
-    LOGE("Failed to select: %s", strerror(errno));
+  ret = send(tcp_socket->fd, buf, len, 0);
+  if (ret < 0) {
+    LOGE("Failed to send: %s", strerror(errno));
     return -1;
   }
-
-  if (FD_ISSET(tcp_socket->fd, &write_set)) {
-    ret = send(tcp_socket->fd, buf, len, 0);
-    if (ret < 0) {
-      LOGE("Failed to send: %s", strerror(errno));
-      return -1;
-    }
-  }
-
   return ret;
 }
 
 int tcp_socket_recv(TcpSocket* tcp_socket, uint8_t* buf, int len) {
-  fd_set read_set;
-  struct timeval tv;
   int ret;
 
   if (tcp_socket->fd < 0) {
@@ -294,23 +276,10 @@ int tcp_socket_recv(TcpSocket* tcp_socket, uint8_t* buf, int len) {
     return -1;
   }
 
-  FD_ZERO(&read_set);
-  FD_SET(tcp_socket->fd, &read_set);
-  tv.tv_sec = 0;
-  tv.tv_usec = 500000;
-
-  if ((ret = select(tcp_socket->fd + 1, &read_set, NULL, NULL, &tv)) < 0) {
-    LOGE("Failed to select: %s", strerror(errno));
+  ret = recv(tcp_socket->fd, buf, len, 0);
+  if (ret < 0) {
+    LOGE("Failed to recv: %s", strerror(errno));
     return -1;
   }
-
-  if (FD_ISSET(tcp_socket->fd, &read_set)) {
-    ret = recv(tcp_socket->fd, buf, len, 0);
-    if (ret < 0) {
-      LOGE("Failed to recv: %s", strerror(errno));
-      return -1;
-    }
-  }
-
   return ret;
 }
diff --git a/src/ssl_transport.c b/src/ssl_transport.c
index f8ebc7a..5d09699 100644
--- a/src/ssl_transport.c
+++ b/src/ssl_transport.c
@@ -13,7 +13,25 @@
 #include "ssl_transport.h"
 #include "utils.h"
 
-static int ssl_transport_mbedtls_recv(void* ctx, unsigned char* buf, size_t len) {
+#define SSL_RECV_TIMEOUT 1000
+
+static int ssl_transport_mbedtls_recv_timeout(void* ctx, unsigned char* buf, size_t len, uint32_t timeout) {
+  int ret;
+  fd_set read_fds;
+  struct timeval tv;
+  tv.tv_sec = timeout / 1000;
+  tv.tv_usec = (timeout % 1000) * 1000;
+
+  FD_ZERO(&read_fds);
+  FD_SET(((TcpSocket*)ctx)->fd, &read_fds);
+
+  ret = select(((TcpSocket*)ctx)->fd + 1, &read_fds, NULL, NULL, &tv);
+  if (ret < 0) {
+    return -1;
+  } else if (ret == 0) {
+    return MBEDTLS_ERR_SSL_TIMEOUT;
+  }
+
   return tcp_socket_recv((TcpSocket*)ctx, buf, len);
 }
 
@@ -78,8 +96,9 @@ int ssl_transport_connect(NetworkContext_t* net_ctx,
     return -1;
   }
 
+  mbedtls_ssl_conf_read_timeout(&net_ctx->conf, SSL_RECV_TIMEOUT);
   mbedtls_ssl_set_bio(&net_ctx->ssl, &net_ctx->tcp_socket,
-                      ssl_transport_mbedlts_send, ssl_transport_mbedtls_recv, NULL);
+                      ssl_transport_mbedlts_send, NULL, ssl_transport_mbedtls_recv_timeout);
 
   LOGI("start to handshake");
 
@@ -90,7 +109,6 @@ int ssl_transport_connect(NetworkContext_t* net_ctx,
   }
 
   LOGI("handshake success");
-
   return 0;
 }