udp.cpp

#include <qthing.h>

#include <map>

#include "esp_log.h"

#include "lwip/err.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include <lwip/netdb.h>

static const char *TAG = "udp-server";

static const in_port_t PORT = 4213;

static std::map<std::string, qthing::udpPacketCallback> packetCallbackMap;

void qthing::addUDPPacketCallback(std::string magicString, qthing::udpPacketCallback callback) {
    if (packetCallbackMap.count(magicString) > 0) {
        qthing::udpPacketCallback old_callback = packetCallbackMap.at(magicString);
        packetCallbackMap[magicString] = [old_callback, callback](udpPacket packet) {
            old_callback(packet);
            callback(packet);
        };
    } else {
        packetCallbackMap[magicString] = callback;
    }

    ESP_LOGI(TAG, "Registered callback for packets with magic string %s", magicString.c_str());
}

static char rx_buffer[2048];  // TODO: get MTU

static void udp_server_task(void *pvParameters)
{
    char addr_str[128];
    int addr_family;
    int ip_protocol;

    struct sockaddr_in6 destAddr;
    bzero(&destAddr.sin6_addr.un, sizeof(destAddr.sin6_addr.un));
    destAddr.sin6_family = AF_INET6;
    destAddr.sin6_port = htons(PORT);
    addr_family = AF_INET6;
    ip_protocol = IPPROTO_IPV6;
    inet6_ntoa_r(destAddr.sin6_addr, addr_str, sizeof(addr_str) - 1);

    while (true) {
        int sock = socket(addr_family, SOCK_DGRAM, ip_protocol);
        if (sock < 0) {
            ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
            break;
        }
        ESP_LOGI(TAG, "Socket created");

        int err = bind(sock, (struct sockaddr *)&destAddr, sizeof(destAddr));
        if (err < 0) {
            ESP_LOGE(TAG, "Socket unable to bind: errno %d", errno);
        }
        ESP_LOGI(TAG, "Socket bound");
        ESP_LOGI(TAG, "Waiting for data");

        while (true) {
            struct sockaddr_in6 sourceAddr; // Large enough for both IPv4 or IPv6
            socklen_t socklen = sizeof(sourceAddr);

            memset(rx_buffer, 0, sizeof(rx_buffer));
            int len = recvfrom(sock, rx_buffer, sizeof(rx_buffer) - 1, 0, (struct sockaddr *)&sourceAddr, &socklen);

            // Error occured during receiving
            if (len < 0) {
                ESP_LOGE(TAG, "recvfrom failed: errno %d", errno);
                break;
            }
            // Data received
            else {
                ESP_LOGI(TAG, "Received %d bytes from %s", len, to_string(sourceAddr).c_str());

                std::string payload_str(rx_buffer);
                int16_t lenBestMatch = -1;
                qthing::udpPacketCallback callback = NULL;
                for(auto it = packetCallbackMap.begin(); it != packetCallbackMap.end(); ++it) {
                    if (len >= it->first.length() && payload_str.find(it->first) == 0) {
                        lenBestMatch = len;
                        callback = it->second;
                    }
                }

                if (lenBestMatch > -1) {
                    qthing::udpPacket packet = {
                        .sourceAddr = sourceAddr,
                        .payload = rx_buffer,
                        .length = (uint16_t)len
                    };

                    callback(packet);
                } else {
                    ESP_LOGW(TAG, "Packet has unknown magic string");
                }
            }
        }

        if (sock != -1) {
            ESP_LOGE(TAG, "Shutting down socket and restarting...");
            shutdown(sock, 0);
            close(sock);
        }
    }
    vTaskDelete(NULL);
}

void start_udp_server() {
    xTaskCreate(udp_server_task, "udp_server", 4096, NULL, 5, NULL);
}