me-no-dev · fremouw · Sep 26, 2018 · Oct 25, 2018 · Oct 25, 2018 · Nov 5, 2018
diff --git a/README.md b/README.md
@@ -9,3 +9,10 @@ This library is the base for [ESPAsyncWebServer](https://github.com/me-no-dev/ES
 
 ## AsyncClient and AsyncServer
 The base classes on which everything else is built. They expose all possible scenarios, but are really raw and require more skills to use.
+
+## TLS support
+Support for TLS is added using mbed TLS, for now only the client part is supported. You can enable this by adding the flag ASYNC_TCP_SSL_ENABLED to your build flags (-DASYNC_TCP_SSL_ENABLED). If you'd like to set a root certificate you can use the setRootCa function on AsyncClient. Feel free to add support for the server side as well :-)
+
+In addition to the regular certificate based cipher suites there is also support for Pre-Shared Key
+cipher suites. Use `setPsk` to define the PSK identifier and PSK itself. The PSK needs to be
+provided in the form of a hex string (and easy way to generate a PSK is to use md5sum).
diff --git a/component.mk b/component.mk
@@ -1,3 +1,5 @@
 COMPONENT_ADD_INCLUDEDIRS := src
 COMPONENT_SRCDIRS := src
 CXXFLAGS += -fno-rtti
+CXXFLAGS += -DASYNC_TCP_SSL_ENABLED=1
+CPPFLAGS += -DASYNC_TCP_SSL_ENABLED=1
diff --git a/src/AsyncTCP.cpp b/src/AsyncTCP.cpp
@@ -408,7 +408,6 @@ static tcp_pcb * _tcp_listen_with_backlog(tcp_pcb * pcb, uint8_t backlog) {
 /*
   Async TCP Client
  */
-
 AsyncClient::AsyncClient(tcp_pcb* pcb)
 : _connect_cb(0)
 , _connect_cb_arg(0)
@@ -425,6 +424,14 @@ AsyncClient::AsyncClient(tcp_pcb* pcb)
 , _timeout_cb(0)
 , _timeout_cb_arg(0)
 , _pcb_busy(false)
+#if ASYNC_TCP_SSL_ENABLED
+, _root_ca_len(0)
+, _root_ca(NULL)
+, _pcb_secure(false)
+, _handshake_done(true)
+, _psk_ident(0)
+, _psk(0)
+#endif // ASYNC_TCP_SSL_ENABLED
 , _pcb_sent_at(0)
 , _close_pcb(false)
 , _ack_pcb(true)
@@ -436,8 +443,6 @@ AsyncClient::AsyncClient(tcp_pcb* pcb)
 , next(NULL)
 , _in_lwip_thread(false)
 {
-    //ets_printf("+: 0x%08x\n", (uint32_t)this);
-
     _pcb = pcb;
     if(_pcb){
         _rx_last_packet = millis();
@@ -453,11 +458,13 @@ AsyncClient::AsyncClient(tcp_pcb* pcb)
 AsyncClient::~AsyncClient(){
     if(_pcb)
         _close();
-
-    //ets_printf("-: 0x%08x\n", (uint32_t)this);
 }
 
+#if ASYNC_TCP_SSL_ENABLED
+bool AsyncClient::connect(IPAddress ip, uint16_t port, bool secure){
+#else
 bool AsyncClient::connect(IPAddress ip, uint16_t port){
+#endif // ASYNC_TCP_SSL_ENABLED
     if (_pcb){
         log_w("already connected, state %d", _pcb->state);
         return false;
@@ -477,6 +484,11 @@ bool AsyncClient::connect(IPAddress ip, uint16_t port){
         return false;
     }
 
+#if ASYNC_TCP_SSL_ENABLED
+    _pcb_secure = secure;
+    _handshake_done = !secure;
+#endif // ASYNC_TCP_SSL_ENABLED
+
     tcp_arg(pcb, this);
     tcp_err(pcb, &_tcp_error);
     if(_in_lwip_thread){
@@ -487,6 +499,18 @@ bool AsyncClient::connect(IPAddress ip, uint16_t port){
     return true;
 }
 
+#if ASYNC_TCP_SSL_ENABLED
+void AsyncClient::setRootCa(const char* rootca, const size_t len) {
+    _root_ca = (char*)rootca;
+    _root_ca_len = len;
+}
+
+void AsyncClient::setPsk(const char* psk_ident, const char* psk) {
+    _psk_ident = psk_ident;
+    _psk = psk;
+}
+#endif // ASYNC_TCP_SSL_ENABLED
+
 AsyncClient& AsyncClient::operator=(const AsyncClient& other){
     if (_pcb)
         _close();
@@ -499,37 +523,80 @@ AsyncClient& AsyncClient::operator=(const AsyncClient& other){
         tcp_sent(_pcb, &_tcp_sent);
         tcp_err(_pcb, &_tcp_error);
         tcp_poll(_pcb, &_tcp_poll, 1);
+
+#if ASYNC_TCP_SSL_ENABLED
+        if(tcp_ssl_has(_pcb)){
+            _pcb_secure = true;
+            _handshake_done = false;
+            tcp_ssl_arg(_pcb, this);
+            tcp_ssl_data(_pcb, &_s_data);
+            tcp_ssl_handshake(_pcb, &_s_handshake);
+            tcp_ssl_err(_pcb, &_s_ssl_error);
+        } else {
+            _pcb_secure = false;
+            _handshake_done = true;
+        }
+#endif // ASYNC_TCP_SSL_ENABLED
     }
     return *this;
 }
 
 int8_t AsyncClient::_connected(void* pcb, int8_t err){
     _pcb = reinterpret_cast<tcp_pcb*>(pcb);
+
     if(_pcb){
         _rx_last_packet = millis();
         _pcb_busy = false;
         tcp_recv(_pcb, &_tcp_recv);
         tcp_sent(_pcb, &_tcp_sent);
         tcp_poll(_pcb, &_tcp_poll, 1);
+#if ASYNC_TCP_SSL_ENABLED
+        if(_pcb_secure){
+            bool err = false;
+            if(_root_ca) {
+              err = tcp_ssl_new_client(_pcb, _hostname.empty() ? NULL : _hostname.c_str(), _root_ca, _root_ca_len) < 0;
+            } else {
+              err = tcp_ssl_new_psk_client(_pcb, _psk_ident, _psk) < 0;
+            }
+            if (err) {
+                log_e("closing....");
+                return _close();
+            }
+
+            tcp_ssl_arg(_pcb, this);
+            tcp_ssl_data(_pcb, &_s_data);
+            tcp_ssl_handshake(_pcb, &_s_handshake);
+            tcp_ssl_err(_pcb, &_s_ssl_error);
+        }
+#endif // ASYNC_TCP_SSL_ENABLED
     }
+
     _in_lwip_thread = true;
+#if ASYNC_TCP_SSL_ENABLED
+    if(!_pcb_secure && _connect_cb)
+#else
     if(_connect_cb)
+#endif // ASYNC_TCP_SSL_ENABLED
         _connect_cb(_connect_cb_arg, this);
     _in_lwip_thread = false;
+
     return ERR_OK;
 }
 
 int8_t AsyncClient::_close(){
-    //ets_printf("X: 0x%08x\n", (uint32_t)this);
     int8_t err = ERR_OK;
+
     if(_pcb) {
-        //log_i("");
+#if ASYNC_TCP_SSL_ENABLED
+        if(_pcb_secure){
+            tcp_ssl_free(_pcb);
+        }
+#endif // ASYNC_TCP_SSL_ENABLED
         tcp_arg(_pcb, NULL);
         tcp_sent(_pcb, NULL);
         tcp_recv(_pcb, NULL);
         tcp_err(_pcb, NULL);
         tcp_poll(_pcb, NULL, 0);
-        _tcp_clear_events(this);
         if(_in_lwip_thread){
             err = tcp_close(_pcb);
         } else {
@@ -539,13 +606,15 @@ int8_t AsyncClient::_close(){
             err = abort();
         }
         _pcb = NULL;
+        // _tcp_clear_events(this);
         if(_discard_cb)
             _discard_cb(_discard_cb_arg, this);
     }
     return err;
 }
 
 void AsyncClient::_error(int8_t err) {
+    log_e("Error!! %d", err);
     if(_pcb){
         tcp_arg(_pcb, NULL);
         tcp_sent(_pcb, NULL);
@@ -560,7 +629,19 @@ void AsyncClient::_error(int8_t err) {
         _discard_cb(_discard_cb_arg, this);
 }
 
+#if ASYNC_TCP_SSL_ENABLED
+void AsyncClient::_ssl_error(int8_t err){
+  if(_error_cb)
+    _error_cb(_error_cb_arg, this, err+64);
+}
+#endif // ASYNC_TCP_SSL_ENABLED
+
 int8_t AsyncClient::_sent(tcp_pcb* pcb, uint16_t len) {
+#if ASYNC_TCP_SSL_ENABLED
+    if (_pcb_secure && !_handshake_done)
+        return ERR_OK;
+#endif // ASYNC_TCP_SSL_ENABLED
+
     _in_lwip_thread = false;
     _rx_last_packet = millis();
     //log_i("%u", len);
@@ -572,19 +653,35 @@ int8_t AsyncClient::_sent(tcp_pcb* pcb, uint16_t len) {
 
 int8_t AsyncClient::_recv(tcp_pcb* pcb, pbuf* pb, int8_t err) {
     if(!_pcb || pcb != _pcb){
-        log_e("0x%08x != 0x%08x", (uint32_t)pcb, (uint32_t)_pcb);
-        if(pb){
-            pbuf_free(pb);
-        }
-        return ERR_OK;
-    }
+         log_e("0x%08x != 0x%08x", (uint32_t)pcb, (uint32_t)_pcb);
+         if(pb){
+             pbuf_free(pb);
+         }
+         return ERR_OK;
+     }
+
     _in_lwip_thread = false;
     if(pb == NULL){
         return _close();
     }
 
     while(pb != NULL){
         _rx_last_packet = millis();
+#if ASYNC_TCP_SSL_ENABLED
+        if(_pcb_secure){
+            // log_i("_recv: %d\n", pb->tot_len);
+            int read_bytes = tcp_ssl_read(pcb, pb);
+            if(read_bytes < 0){
+                if (read_bytes != MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
+                    log_e("_recv err: %d\n", read_bytes);
+                    _close();
+                }
+
+                //return read_bytes;
+            }
+            return ERR_OK;
+        }
+#endif // ASYNC_TCP_SSL_ENABLED
         //we should not ack before we assimilate the data
         //log_i("%u", pb->len);
         //Serial.write((const uint8_t *)pb->payload, pb->len);
@@ -631,6 +728,12 @@ int8_t AsyncClient::_poll(tcp_pcb* pcb){
         _close();
         return ERR_OK;
     }
+#if ASYNC_TCP_SSL_ENABLED
+    if(_pcb_secure && !_handshake_done && (now - _rx_last_packet) >= 2000){
+        _close();
+        return ERR_OK;
+    }
+#endif // ASYNC_TCP_SSL_ENABLED
     // Everything is fine
     if(_poll_cb)
         _poll_cb(_poll_cb_arg, this);
@@ -639,8 +742,13 @@ int8_t AsyncClient::_poll(tcp_pcb* pcb){
 
 void AsyncClient::_dns_found(struct ip_addr *ipaddr){
     _in_lwip_thread = true;
+
     if(ipaddr){
+#if ASYNC_TCP_SSL_ENABLED
+        connect(IPAddress(ipaddr->u_addr.ip4.addr), _connect_port, _pcb_secure);
+#else
         connect(IPAddress(ipaddr->u_addr.ip4.addr), _connect_port);
+#endif // ASYNC_TCP_SSL_ENABLED
     } else {
         log_e("dns fail");
         if(_error_cb)
@@ -655,13 +763,29 @@ bool AsyncClient::operator==(const AsyncClient &other) {
     return _pcb == other._pcb;
 }
 
+#if ASYNC_TCP_SSL_ENABLED
+bool AsyncClient::connect(const char* host, uint16_t port, bool secure){
+#else
 bool AsyncClient::connect(const char* host, uint16_t port){
+#endif // ASYNC_TCP_SSL_ENABLED
     ip_addr_t addr;
+
     err_t err = dns_gethostbyname(host, &addr, (dns_found_callback)&_s_dns_found, this);
     if(err == ERR_OK) {
+        _hostname = host;
+
+#if ASYNC_TCP_SSL_ENABLED
+        return connect(IPAddress(addr.u_addr.ip4.addr), port, secure);
+#else
         return connect(IPAddress(addr.u_addr.ip4.addr), port);
+#endif // ASYNC_TCP_SSL_ENABLED
     } else if(err == ERR_INPROGRESS) {
+        _hostname = host;
         _connect_port = port;
+#if ASYNC_TCP_SSL_ENABLED
+        _pcb_secure = secure;
+        _handshake_done = !secure;
+#endif // ASYNC_TCP_SSL_ENABLED
         return true;
     }
     log_e("error: %d", err);
@@ -727,13 +851,25 @@ size_t AsyncClient::write(const char* data, size_t size, uint8_t apiflags) {
     return will_send;
 }
 
-
 size_t AsyncClient::add(const char* data, size_t size, uint8_t apiflags) {
     if(!_pcb || size == 0 || data == NULL)
         return 0;
     size_t room = space();
     if(!room)
         return 0;
+#if ASYNC_TCP_SSL_ENABLED
+    if(_pcb_secure){
+        int sent = tcp_ssl_write(_pcb, (uint8_t*)data, size);
+        if(sent >= 0){
+            // @ToDo: ???
+            //_tx_unacked_len += sent;
+            return sent;
+        }
+        //log_i("add: tcp_ssl_write: %d", sent);
+        _close();
+        return 0;
+    }
+#endif // ASYNC_TCP_SSL_ENABLED
     size_t will_send = (room < size) ? room : size;
     int8_t err = ERR_OK;
     if(_in_lwip_thread){
@@ -747,6 +883,10 @@ size_t AsyncClient::add(const char* data, size_t size, uint8_t apiflags) {
 }
 
 bool AsyncClient::send(){
+#if ASYNC_TCP_SSL_ENABLED
+    if(_pcb_secure)
+        return true;
+#endif // ASYNC_TCP_SSL_ENABLED
     int8_t err = ERR_OK;
     if(_in_lwip_thread){
         err = tcp_output(_pcb);
@@ -958,7 +1098,6 @@ void AsyncClient::onPoll(AcConnectHandler cb, void* arg){
     _poll_cb_arg = arg;
 }
 
-
 void AsyncClient::_s_dns_found(const char * name, struct ip_addr * ipaddr, void * arg){
     if(arg){
         reinterpret_cast<AsyncClient*>(arg)->_dns_found(ipaddr);
@@ -981,7 +1120,7 @@ int8_t AsyncClient::_s_recv(void * arg, struct tcp_pcb * pcb, struct pbuf *pb, i
         reinterpret_cast<AsyncClient*>(arg)->_recv(pcb, pb, err);
     } else {
         if(pb){
-           pbuf_free(pb);
+        pbuf_free(pb);
         }
         log_e("Bad Args: 0x%08x 0x%08x", arg, pcb);
     }
@@ -1014,6 +1153,25 @@ int8_t AsyncClient::_s_connected(void * arg, void * pcb, int8_t err){
     return ERR_OK;
 }
 
+#if ASYNC_TCP_SSL_ENABLED
+void AsyncClient::_s_data(void *arg, struct tcp_pcb *tcp, uint8_t * data, size_t len){
+  AsyncClient *c = reinterpret_cast<AsyncClient*>(arg);
+  if(c->_recv_cb)
+    c->_recv_cb(c->_recv_cb_arg, c, data, len);
+}
+
+void AsyncClient::_s_handshake(void *arg, struct tcp_pcb *tcp, struct tcp_ssl_pcb* ssl){
+  AsyncClient *c = reinterpret_cast<AsyncClient*>(arg);
+  c->_handshake_done = true;
+  if(c->_connect_cb)
+    c->_connect_cb(c->_connect_cb_arg, c);
+}
+
+void AsyncClient::_s_ssl_error(void *arg, struct tcp_pcb *tcp, int8_t err){
+  reinterpret_cast<AsyncClient*>(arg)->_ssl_error(err);
+}
+#endif // ASYNC_TCP_SSL_ENABLED
+
 const char * AsyncClient::errorToString(int8_t error){
     switch(error){
         case 0: return "OK";