UDPStream.h

#pragma once
#include <WiFi.h>
#include <WiFiUdp.h>
#if defined(ESP32)
#include <esp_wifi.h>
#endif
#include "ArduinoLogger.h"
 
namespace arduino {
 
class UDPStream : public Stream {
 public:
  UDPStream() = default;
 
  UDPStream(IPAddress a, uint16_t port) { begin(a, port); }
 
  void setUDP(UDP& udp) { p_udp = &udp; };
 
  int availableForWrite() override { return 1492; }
 
  int available() override {
    int size = p_udp->available();
    // if the curren package is used up we prvide the info for the next
    if (size == 0) {
      size = p_udp->parsePacket();
    }
    return size;
  }
 
  bool begin(IPAddress a, uint16_t port) {
    remote_address_ext = a;
    remote_port_ext = port;
    return p_udp->begin(port);
  }
 
  bool begin(uint16_t port, uint16_t port_ext = 0) {
    remote_address_ext = IPAddress((uint32_t)0);
    remote_port_ext = port_ext != 0 ? port_ext : port;
    return p_udp->begin(port);
  }
 
  bool beginMulticast(IPAddress address, uint16_t port) {
    return p_udp->beginMulticast(address, port);
  }
 
  bool setTarget(IPAddress a, uint16_t port) { return begin(a, port); }
 
  uint16_t remotePort() {
    uint16_t result = p_udp->remotePort();
    return result != 0 ? result : remote_port_ext;
  }
 
  IPAddress remoteIP() {
    // Determine address if it has not been specified
    if ((uint32_t)remote_address_ext == 0) {
      remote_address_ext = p_udp->remoteIP();
    }
    // EmulatorLogger.log(EmulatorLogger::INFO, "ip: %u", remote_address_ext);
    return remote_address_ext;
  }
 
  size_t write(const uint8_t* data, size_t len) override {
    IPAddress remote = remoteIP();
    if ((uint32_t)remote == 0) {
      Logger.error("UDPStream", "no remote address defined");
      return 0;
    }
    p_udp->beginPacket(remote, remotePort());
    size_t result = p_udp->write(data, len);
    p_udp->endPacket();
    return result;
  }
 
  size_t readBytes(uint8_t* data, size_t len) {
    size_t avail = available();
    size_t bytes_read = 0;
    if (avail > 0) {
      // get the data now
      bytes_read = p_udp->readBytes((uint8_t*)data, len);
    }
    return bytes_read;
  }
 
  void stop() { p_udp->stop(); }
 
  int peek() override { return p_udp->peek(); }
 
  int read() override { return p_udp->read(); }
 
  size_t write(uint8_t data) override {
    if (write_buffer.capacity() < write_buffer_size) {
      write_buffer.reserve(write_buffer_size);
    }
    write_buffer.push_back(data);
    if (write_buffer.size() >= write_buffer_size) {
      flush();
    }
    return 1;
  }
 
  void flush() override {
    if (write_buffer.empty()) {
      return;
    }
    write(write_buffer.data(), write_buffer.size());
    write_buffer.clear();
  }
 
  void setWriteBufferSize(size_t size) { write_buffer_size = size; }
 
 protected:
  WiFiUDP default_udp;
  UDP* p_udp = &default_udp;
  uint16_t remote_port_ext = 0;
  IPAddress remote_address_ext;
  int write_buffer_size = 256;
  std::vector<uint8_t> write_buffer;
};
 
}  // namespace arduino
 
using WiFiUDPStream = arduino::UDPStream;