Add connect timeout to WiFiClient (#2383)

* Add timeout to WiFiClient.connect() * Changed default handling
espressif · Feb 14, 2019 · fc737e0 · fc737e0
1 parent e302a68
commit fc737e0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/libraries/WiFi/src/WiFiClient.cpp b/libraries/WiFi/src/WiFiClient.cpp
@@ -202,38 +202,56 @@ 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;
     bzero((char *) &serveraddr, sizeof(serveraddr));
     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;
     return 1;
 }
 
 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
@@ -41,7 +41,9 @@ class WiFiClient : public Client
     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);