From fc737e08c617d27414ebf1a574f7ebdf9a968f38 Mon Sep 17 00:00:00 2001
From: boarchuz <46267286+boarchuz@users.noreply.github.com>
Date: Thu, 14 Feb 2019 22:55:50 +1100
Subject: [PATCH] Add connect timeout to WiFiClient (#2383)

* Add timeout to WiFiClient.connect()

* Changed default handling
---
 libraries/WiFi/src/WiFiClient.cpp | 26 ++++++++++++++++++++++----
 libraries/WiFi/src/WiFiClient.h   |  2 ++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/libraries/WiFi/src/WiFiClient.cpp b/libraries/WiFi/src/WiFiClient.cpp
index 8ea6c7af..359e90b8 100644
--- a/libraries/WiFi/src/WiFiClient.cpp
+++ b/libraries/WiFi/src/WiFiClient.cpp
@@ -202,12 +202,17 @@ void WiFiClient::stop()
 }
 
 int WiFiClient::connect(IPAddress ip, uint16_t port)
+{
+    return connect(ip,port,-1);
+}
+int WiFiClient::connect(IPAddress ip, uint16_t port, int32_t timeout)
 {
     int sockfd = socket(AF_INET, SOCK_STREAM, 0);
     if (sockfd < 0) {
         log_e("socket: %d", errno);
         return 0;
     }
+    fcntl( sockfd, F_SETFL, fcntl( sockfd, F_GETFL, 0 ) | O_NONBLOCK );
 
     uint32_t ip_addr = ip;
     struct sockaddr_in serveraddr;
@@ -215,12 +220,21 @@ int WiFiClient::connect(IPAddress ip, uint16_t port)
     serveraddr.sin_family = AF_INET;
     bcopy((const void *)(&ip_addr), (void *)&serveraddr.sin_addr.s_addr, 4);
     serveraddr.sin_port = htons(port);
-    int res = lwip_connect_r(sockfd, (struct sockaddr*)&serveraddr, sizeof(serveraddr));
-    if (res < 0) {
-        log_e("lwip_connect_r: %d", errno);
+    fd_set fdset;
+    struct timeval tv;
+    FD_ZERO(&fdset);
+    FD_SET(sockfd, &fdset);
+    tv.tv_sec = 0;
+    tv.tv_usec = timeout * 1000;
+    lwip_connect_r(sockfd, (struct sockaddr*)&serveraddr, sizeof(serveraddr));
+    int res = select(sockfd + 1, nullptr, &fdset, nullptr, timeout<0 ? nullptr : &tv);
+    if (res != 1)
+    {
+        log_e("select: %d",errno);
         close(sockfd);
         return 0;
     }
+    fcntl( sockfd, F_SETFL, fcntl( sockfd, F_GETFL, 0 ) & (~O_NONBLOCK) );
     clientSocketHandle.reset(new WiFiClientSocketHandle(sockfd));
     _rxBuffer.reset(new WiFiClientRxBuffer(sockfd));
     _connected = true;
@@ -228,12 +242,16 @@ int WiFiClient::connect(IPAddress ip, uint16_t port)
 }
 
 int WiFiClient::connect(const char *host, uint16_t port)
+{
+    return connect(host,port,-1);
+}
+int WiFiClient::connect(const char *host, uint16_t port, int32_t timeout)
 {
     IPAddress srv((uint32_t)0);
     if(!WiFiGenericClass::hostByName(host, srv)){
         return 0;
     }
-    return connect(srv, port);
+    return connect(srv, port, timeout);
 }
 
 int WiFiClient::setSocketOption(int option, char* value, size_t len)
diff --git a/libraries/WiFi/src/WiFiClient.h b/libraries/WiFi/src/WiFiClient.h
index e87708f5..09841f1a 100644
--- a/libraries/WiFi/src/WiFiClient.h
+++ b/libraries/WiFi/src/WiFiClient.h
@@ -41,7 +41,9 @@ public:
     WiFiClient(int fd);
     ~WiFiClient();
     int connect(IPAddress ip, uint16_t port);
+    int connect(IPAddress ip, uint16_t port, int32_t timeout);
     int connect(const char *host, uint16_t port);
+    int connect(const char *host, uint16_t port, int32_t timeout);
     size_t write(uint8_t data);
     size_t write(const uint8_t *buf, size_t size);
     size_t write_P(PGM_P buf, size_t size);