From eebcc656edcfae35494c7c11481cd377e96aac9d Mon Sep 17 00:00:00 2001
From: Ivan Grokhotkov <ivan@espressif.com>
Date: Thu, 18 May 2017 17:24:57 +0800
Subject: [PATCH] WiFiClient: add support for connect timeout

---
 libraries/ESP8266WiFi/src/WiFiClient.cpp      | 48 +++----------
 .../ESP8266WiFi/src/include/ClientContext.h   | 70 +++++++++++++++----
 2 files changed, 65 insertions(+), 53 deletions(-)

diff --git a/libraries/ESP8266WiFi/src/WiFiClient.cpp b/libraries/ESP8266WiFi/src/WiFiClient.cpp
index 144cf194b..8f9eada72 100644
--- a/libraries/ESP8266WiFi/src/WiFiClient.cpp
+++ b/libraries/ESP8266WiFi/src/WiFiClient.cpp
@@ -88,7 +88,7 @@ WiFiClient& WiFiClient::operator=(const WiFiClient& other)
 int WiFiClient::connect(const char* host, uint16_t port)
 {
     IPAddress remote_addr;
-    if (WiFi.hostByName(host, remote_addr))
+    if (WiFi.hostByName(host, remote_addr, _timeout))
     {
         return connect(remote_addr, port);
     }
@@ -122,37 +122,19 @@ int WiFiClient::connect(IPAddress ip, uint16_t port)
         pcb->local_port = _localPort++;
     }
 
-    tcp_arg(pcb, this);
-    tcp_err(pcb, &WiFiClient::_s_err);
-    tcp_connect(pcb, &addr, port, reinterpret_cast<tcp_connected_fn>(&WiFiClient::_s_connected));
-
-    esp_yield();
-    if (_client)
-        return 1;
-
-    //  if tcp_error was called, pcb has already been destroyed.
-    // tcp_abort(pcb);
-    return 0;
-}
-
-int8_t WiFiClient::_connected(void* pcb, int8_t err)
-{
-    (void) err;
-    tcp_pcb* tpcb = reinterpret_cast<tcp_pcb*>(pcb);
-    _client = new ClientContext(tpcb, 0, 0);
+    _client = new ClientContext(pcb, nullptr, nullptr);
     _client->ref();
-    esp_schedule();
-    return ERR_OK;
-}
+    _client->setTimeout(_timeout);
+    int res = _client->connect(&addr, port);
+    if (res == 0) {
+        _client->unref();
+        _client = nullptr;
+        return 0;
+    }
 
-void WiFiClient::_err(int8_t err)
-{
-    (void) err;
-    DEBUGV(":err %d\r\n", err);
-    esp_schedule();
+    return 1;
 }
 
-
 void WiFiClient::setNoDelay(bool nodelay) {
     if (!_client)
         return;
@@ -331,16 +313,6 @@ uint16_t WiFiClient::localPort()
     return _client->getLocalPort();
 }
 
-int8_t WiFiClient::_s_connected(void* arg, void* tpcb, int8_t err)
-{
-    return reinterpret_cast<WiFiClient*>(arg)->_connected(tpcb, err);
-}
-
-void WiFiClient::_s_err(void* arg, int8_t err)
-{
-    reinterpret_cast<WiFiClient*>(arg)->_err(err);
-}
-
 void WiFiClient::stopAll()
 {
     for (WiFiClient* it = _s_first; it; it = it->_next) {
diff --git a/libraries/ESP8266WiFi/src/include/ClientContext.h b/libraries/ESP8266WiFi/src/include/ClientContext.h
index c7d5bde00..82e059467 100644
--- a/libraries/ESP8266WiFi/src/include/ClientContext.h
+++ b/libraries/ESP8266WiFi/src/include/ClientContext.h
@@ -124,7 +124,25 @@ public:
         }
     }
 
-    size_t availableForWrite ()
+    int connect(ip_addr_t* addr, uint16_t port)
+    {
+        err_t err = tcp_connect(_pcb, addr, port, reinterpret_cast<tcp_connected_fn>(&ClientContext::_s_connected));
+        if (err != ERR_OK) {
+            return 0;
+        }
+        _connect_pending = 1;
+        _op_start_time = millis();
+        // This delay will be interrupted by esp_schedule in the connect callback
+        delay(_timeout_ms);
+        _connect_pending = 0;
+        if (state() != ESTABLISHED) {
+            abort();
+            return 0;
+        }
+        return 1;
+    }
+
+    size_t availableForWrite()
     {
         return _pcb? tcp_sndbuf(_pcb): 0;
     }
@@ -149,14 +167,14 @@ public:
         return tcp_nagle_disabled(_pcb);
     }
 
-    void setNonBlocking(bool nonblocking)
+    void setTimeout(int timeout_ms) 
     {
-        _noblock = nonblocking;
+        _timeout_ms = timeout_ms;
     }
 
-    bool getNonBlocking()
+    int getTimeout()
     {
-        return _noblock;
+        return _timeout_ms;
     }
 
     uint32_t getRemoteAddress()
@@ -315,9 +333,14 @@ public:
 
 protected:
 
-    void _cancel_write()
+    bool _is_timeout()
     {
-        if (_send_waiting) {
+        return millis() - _op_start_time > _timeout_ms;
+    }
+
+    void _notify_error()
+    {
+        if (_connect_pending || _send_waiting) {
             esp_schedule();
         }
     }
@@ -328,10 +351,11 @@ protected:
         assert(_send_waiting == 0);
         _datasource = ds;
         _written = 0;
+        _op_start_time = millis();
         do {
             _write_some();
 
-            if (!_datasource->available() || _noblock || state() == CLOSED) {
+            if (!_datasource->available() || _is_timeout() || state() == CLOSED) {
                 delete _datasource;
                 _datasource = nullptr;
                 break;
@@ -431,7 +455,7 @@ protected:
         (void) err;
         if(pb == 0) { // connection closed
             DEBUGV(":rcl\r\n");
-            _cancel_write();
+            _notify_error();
             abort();
             return ERR_ABRT;
         }
@@ -456,7 +480,16 @@ protected:
         tcp_recv(_pcb, NULL);
         tcp_err(_pcb, NULL);
         _pcb = NULL;
-        _cancel_write();
+        _notify_error();
+    }
+
+    int8_t _connected(void* pcb, int8_t err)
+    {
+        (void) err;
+        assert(pcb == _pcb);
+        assert(_connect_pending);
+        esp_schedule();
+        return ERR_OK;
     }
 
     err_t _poll(tcp_pcb*)
@@ -485,6 +518,11 @@ protected:
         return reinterpret_cast<ClientContext*>(arg)->_sent(tpcb, len);
     }
 
+    static int8_t _s_connected(void* arg, void* pcb, int8_t err)
+    {
+        return reinterpret_cast<ClientContext*>(arg)->_connected(pcb, err);
+    }
+
 private:
     tcp_pcb* _pcb;
 
@@ -494,14 +532,16 @@ private:
     discard_cb_t _discard_cb;
     void* _discard_cb_arg;
 
-    int _refcnt;
-    ClientContext* _next;
-
     DataSource* _datasource = nullptr;
     size_t _written = 0;
     size_t _write_chunk_size = 256;
-    bool _noblock = false;
-    int _send_waiting = 0;
+    uint32_t _timeout_ms = 5000;
+    uint32_t _op_start_time = 0;
+    uint8_t _send_waiting = 0;
+    uint8_t _connect_pending = 0;
+
+    int8_t _refcnt;
+    ClientContext* _next;
 };
 
 #endif//CLIENTCONTEXT_H