UDPHelper_POSIX.h

#pragma once
 
#include <stdio.h>
#include <errno.h>
 
#ifdef _WIN32
#include <winsock2.h>
#include <WS2tcpip.h>
 
#define close(fd) closesocket(fd)
#define ssize_t int
#else
#include <sys/time.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#endif
 
#ifdef __ZEPHYR__
#include <fcntl.h>
#define INADDR_BROADCAST ((u32_t) 0xffffffff)
#endif
 
class UDPHelper
{
    uint16_t _port;
    uint32_t _magic_header;
    sockaddr_in _localaddr, _remote_receiver_addr, _remote_sender_addr;
    int _fd = -1;
public:
    ~UDPHelper()
    {
        if (_fd != -1)
#ifdef _WIN32
            closesocket(_fd);
        WSACleanup();
#else
            close(_fd);
#endif
    }
 
    bool begin(uint16_t port)
    {
#ifdef _WIN32
        // Initialize Winsock
        WSAData wsaData;
        int iResult;
        iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
        if (iResult != 0) {
            return false;
        }
#endif
        _port = port;
 
        // Close if open after previous init attempt
        if (_fd != -1) {
#ifdef _WIN32
            closesocket(_fd);
#else
            close(_fd);
#endif
            _fd = -1;
        }
 
        // Prepare socket
        _fd=socket(AF_INET,SOCK_DGRAM, IPPROTO_UDP);
        if (_fd==-1) {
            //printf("INIT listening socket %s\n", strerror(errno));
            return false;
        }
 
        // Make it non-blocking
#ifdef _WIN32
        DWORD read_timeout = 10; // millis
#else
        struct timeval read_timeout;
        read_timeout.tv_sec = 0;
        read_timeout.tv_usec = 1000;
#endif
#ifdef __ZEPHYR__
        fcntl(_fd, F_SETFL, O_NONBLOCK);
#else
        setsockopt(_fd, SOL_SOCKET, SO_RCVTIMEO, (char *)&read_timeout, sizeof read_timeout);
#endif
 
        // Bind to specific local port
        memset(&_localaddr, 0, sizeof(_localaddr));
        _localaddr.sin_family = AF_INET;
        _localaddr.sin_port = htons(_port);
        _localaddr.sin_addr.s_addr = INADDR_ANY;
        if (bind(_fd,(struct sockaddr *) &_localaddr,sizeof(_localaddr))==-1) {
            //printf("INIT listening bind %s\n", strerror(errno));
            return false;
        }
        memset(&_remote_sender_addr, 0, sizeof(_remote_sender_addr));
 
        // Prepare receiver address as broadcast
        memset(&_remote_receiver_addr, 0, sizeof(_remote_receiver_addr));
        _remote_receiver_addr.sin_family = AF_INET;
        _remote_receiver_addr.sin_port = htons(_port);
        _remote_receiver_addr.sin_addr.s_addr = INADDR_BROADCAST;
 
#ifndef __ZEPHYR__
        // Allow broadcasts
        int broadcast=1;
        if (setsockopt(_fd,SOL_SOCKET,SO_BROADCAST,(const char*)&broadcast,sizeof(broadcast))==-1) {
            //printf("INIT send setsockopt %s\n", strerror(errno));
            return false;
        }
#endif
        return true;
    }
 
    uint16_t receive_frame(uint8_t *string, uint16_t max_length)
    {
        struct sockaddr_storage src_addr;
        socklen_t src_addr_len=sizeof(src_addr);
        ssize_t count=recvfrom(_fd,(char*)string,max_length,0,(struct sockaddr*)&src_addr,&src_addr_len);
        if (count==-1) {
#ifdef _WIN32
            //int error = WSAGetLastError();
            //if (error != WSAETIMEDOUT) printf("recvfrom error %d, bufsize=%d\n", error, max_length);
#endif
            return false; // Reception failed
        } else if (count==max_length || count < 4) {
            //printf("FAIL receive_frame recvfrom %d\n", count);
            return false; // Too large packet
        } else {
            //printf("OK receive_frame recvfrom %d, maxize %d\n", count, max_length);
            // Remember sender's address
            memcpy(&_remote_sender_addr, &src_addr, sizeof(_remote_sender_addr));
 
            // Get header
            uint32_t header = 0;
            memcpy(&header, string, 4);
            if(header != _magic_header) {
                return false;    // Not a LocalUDP packet
            }
 
            // Shift contents to remove header
            for (uint16_t i=0; i<count-4; i++) {
                string[i] = string[i+4];
            }
            return count - 4;
        }
        return PJON_FAIL;
    }
 
    class Buf
    {
        char *buf = NULL;
        uint16_t len = 0;
    public:
        Buf(uint16_t size)
        {
            len = size;
            buf = new char[size];
        }
        ~Buf()
        {
            if (buf) {
                delete[] buf;
            }
        }
        char* operator()()
        {
            return buf;
        }
        uint16_t size()
        {
            return len;
        }
    };
 
 
    void send_frame(const uint8_t *string, uint16_t length, const sockaddr_in &remote_addr)
    {
        if(length > 0) {
            Buf buffer(4 + length);
            memcpy(buffer(), &_magic_header, 4);
            memcpy(&(buffer()[4]), string, length);
            int res = sendto(_fd,buffer(),buffer.size(),0,(const sockaddr *)&remote_addr,sizeof(remote_addr));
            //printf("send_frame %d sendto %d\n", length, res);
        }
    }
 
    void send_response(uint8_t *string, uint16_t length)
    {
        send_frame((const uint8_t *)string, length, _remote_sender_addr);
    }
 
    void send_response(uint8_t response)
    {
        send_frame((const uint8_t *)&response, 1, _remote_sender_addr);
    }
 
    void send_frame(const uint8_t *string, uint16_t length)
    {
        _remote_receiver_addr.sin_port = htons(_port);
        _remote_receiver_addr.sin_addr.s_addr = INADDR_BROADCAST;
        send_frame(string, length, _remote_receiver_addr);
    }
 
    void send_frame(const uint8_t *string, uint16_t length, uint8_t *remote_ip, uint16_t remote_port)
    {
        _remote_receiver_addr.sin_port = htons(remote_port);
        _remote_receiver_addr.sin_addr.s_addr = *(uint32_t*)remote_ip;
        send_frame(string, length, _remote_receiver_addr);
    }
 
    void set_magic_header(uint32_t magic_header)
    {
        _magic_header = magic_header;
    }
 
    void get_sender(uint8_t *ip, uint16_t &port)
    {
        memcpy(ip, &_remote_sender_addr.sin_addr.s_addr, 4);
        port = ntohs(_remote_sender_addr.sin_port);
    }
};
 
#undef close