1 files changed, 42 insertions, 33 deletions

diff --git a/fw/esp32/components/eos/wifi.c b/fw/esp32/components/eos/wifi.c
index 77f9143..fe1b382 100755
--- a/fw/esp32/components/eos/wifi.c
+++ b/fw/esp32/components/eos/wifi.c
@@ -13,6 +13,7 @@
 
 #include "eos.h"
 #include "net.h"
+#include "tun.h"
 #include "wifi.h"
 
 // XXX: WiFi fail due to no DHCP server
@@ -53,6 +54,7 @@ static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t e
 
                 if (sta_disconnected->reason == WIFI_REASON_ASSOC_LEAVE) {
                     eos_wifi_send_status();
+                    eos_tun_portmap_remove();
                     break;
                 }
                 if ((reconnect_cnt == 0) && (sta_disconnected->reason == WIFI_REASON_BEACON_TIMEOUT)) {
@@ -64,6 +66,7 @@ static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t e
                     reconnect_cnt++;
                 } else {
                     eos_wifi_send_status();
+                    eos_tun_portmap_remove();
 
                     reconnect_cnt = 0;
                     if (stop) {
@@ -84,13 +87,17 @@ static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t e
 
                 got_ip = (ip_event_got_ip_t *)event_data;
                 ESP_LOGI(TAG, "Event got IP - addres: " IPSTR, IP2STR(&got_ip->ip_info.ip));
+                ESP_LOGI(TAG, "Event got IP - addres: %x", (unsigned int)got_ip->ip_info.ip.addr);
                 eos_wifi_send_status();
+                eos_tun_portmap_remove();
+                eos_tun_portmap_add(got_ip->ip_info.ip.addr);
 
                 /*
                 if (got_ip->ip_changed) {
                     recreate all sockets
                 }
                 */
+
                 break;
             }
 
@@ -101,35 +108,25 @@ static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t e
 
 }
 
-static void wifi_handler(unsigned char _mtype, unsigned char *buffer, uint16_t buf_len) {
+static void wifi_handler(unsigned char _mtype, EOSMessage *msg, uint16_t len) {
+    unsigned char *buffer = msg->buffer;
+    uint16_t buf_size = msg->size;
     esp_err_t ret = ESP_OK;
-    uint8_t mtype;
+    unsigned char mtype;
 
-    if (buf_len < 1) return;
+    if (len < 1) return;
 
     mtype = buffer[0];
     switch (mtype) {
         case EOS_WIFI_MTYPE_STATUS: {
-            unsigned char *_buffer;
-            ssize_t _rv;
-            int reply, rv;
-
-            reply = _mtype & EOS_NET_MTYPE_FLAG_REPL;
+            ssize_t rv;
 
-            if (reply) {
-                _rv = eos_wifi_get_status(buffer + 1);
-                if (_rv < 0) break;
+            if (!(eos_msg_flags(msg) & EOS_MSG_FLAG_RPLY_REQ)) break;
 
-                eos_net_reply(EOS_NET_MTYPE_WIFI, buffer, _rv + 1);
-            } else {
-                _buffer = eos_net_alloc();
-                _buffer[0] = EOS_WIFI_MTYPE_STATUS;
-                _rv = eos_wifi_get_status(_buffer + 1);
-                if (_rv < 0) break;
+            rv = eos_wifi_get_status(buffer + 1, buf_size - 1);
+            if (rv < 0) break;
 
-                rv = eos_net_send(EOS_NET_MTYPE_WIFI, _buffer, _rv + 1);
-                if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
-            }
+            eos_net_reply(EOS_NET_MTYPE_WIFI, msg, rv + 1);
             break;
         }
 
@@ -233,11 +230,14 @@ void eos_wifi_run(void) {
     ESP_LOGI(TAG, "RUN");
 }
 
-ssize_t eos_wifi_get_status(unsigned char *buffer) {
+ssize_t eos_wifi_get_status(unsigned char *buffer, size_t size) {
     unsigned char *p;
     wifi_ap_record_t ap_info;
     esp_err_t ret;
 
+    /* 1 + sizeof(uint32_t) + sizeof(ap_info.ssid) - 1 */
+    if (size < sizeof(uint32_t) + sizeof(ap_info.ssid)) return EOS_ERR_SIZE;
+
     p = buffer;
     ret = esp_wifi_sta_get_ap_info(&ap_info);
     switch (ret) {
@@ -296,28 +296,33 @@ ssize_t eos_wifi_get_status(unsigned char *buffer) {
 }
 
 void eos_wifi_send_status(void) {
-    unsigned char *rbuf;
+    EOSMessage msg;
     ssize_t _rv;
     int rv;
 
-    rbuf = eos_net_alloc();
-    rbuf[0] = EOS_WIFI_MTYPE_STATUS;
-    _rv = eos_wifi_get_status(rbuf + 1);
+    eos_net_alloc(&msg);
+    if (msg.size < 1) {
+        eos_net_free(&msg);
+        return;
+    }
+    msg.buffer[0] = EOS_WIFI_MTYPE_STATUS;
+    _rv = eos_wifi_get_status(msg.buffer + 1, msg.size - 1);
     if (_rv < 0) {
-        eos_net_free(rbuf);
+        eos_net_free(&msg);
         return;
     }
 
-    rv = eos_net_send(EOS_NET_MTYPE_WIFI, rbuf, _rv + 1);
+    rv = eos_net_send(EOS_NET_MTYPE_WIFI, &msg, _rv + 1);
     if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
 }
 
 void eos_wifi_send_scan(void) {
+    EOSMessage msg;
     static wifi_ap_record_t scan_r[EOS_WIFI_MAX_SCAN_RECORDS];
     static uint16_t scan_n;
     size_t len;
     esp_err_t ret;
-    unsigned char *rbuf, *p;
+    unsigned char *p;
     int i, rv;
 
     scan_n = EOS_WIFI_MAX_SCAN_RECORDS;
@@ -330,19 +335,23 @@ void eos_wifi_send_scan(void) {
 
     ESP_LOGI(TAG, "Scan done: %d", scan_n);
 
-    rbuf = eos_net_alloc();
-    rbuf[0] = EOS_WIFI_MTYPE_SCAN;
-    p = rbuf + 1;
+    eos_net_alloc(&msg);
+    if (msg.size < 1) {
+        eos_net_free(&msg);
+        return;
+    }
+    msg.buffer[0] = EOS_WIFI_MTYPE_SCAN;
+    p = msg.buffer + 1;
 
     for (i=0; i<scan_n; i++) {
         len = strnlen((char *)scan_r[i].ssid, sizeof(scan_r[i].ssid));
         if (len == sizeof(scan_r[i].ssid)) continue;
-        if (p - rbuf + len + 1 > EOS_NET_SIZE_BUF) break;
+        if (p - msg.buffer + len + 1 > msg.size) break;
 
         strcpy((char *)p, (char *)scan_r[i].ssid);
         p += len + 1;
     }
 
-    rv = eos_net_send(EOS_NET_MTYPE_WIFI, rbuf, p - rbuf);
+    rv = eos_net_send(EOS_NET_MTYPE_WIFI, &msg, p - msg.buffer);
     if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
 }