CPP_Core/Core/socket/UDP_Server.cpp

#include "UDP_Server.h"

#include <algorithm>
#include <string_view>

namespace Psc::asio_socket {

void UDP_Server::create() {
  socket = std::make_unique<asio::ip::udp::socket>(io_context);
  socket_fd = socket_id(socket.get());
  recv_buffer.init(buffer_size);
  send_buffer.init(buffer_size);
  recv_storage.assign(static_cast<size_t>(read_chunk_size), 0);
  send_storage.assign(static_cast<size_t>(read_chunk_size), 0);
  read_pending = false;
  write_pending = false;
  state = Not_Bind;
}

bool UDP_Server::bind() {
  if (!socket)
    create();
  if (bind_address.ip.empty() || bind_address.port == 0) {
    state = Not_Set_Field;
    return false;
  }

  asio::error_code ec;
  auto address_value = asio::ip::make_address(bind_address.ip, ec);
  if (ec) {
    state = Not_Bind;
    return false;
  }

  asio::ip::udp::endpoint endpoint(
      address_value, static_cast<unsigned short>(bind_address.port));
  socket->open(endpoint.protocol(), ec);
  if (ec)
    return false;
  socket->set_option(asio::socket_base::reuse_address(true), ec);
  socket->set_option(
      asio::socket_base::receive_buffer_size(static_cast<int>(buffer_size)),
      ec);
  socket->bind(endpoint, ec);
  if (ec) {
    state = Not_Bind;
    return false;
  }
  bound = true;
  state = Working;
  start_read();
  return true;
}

void UDP_Server::tick() {
  if (!bound || !socket)
    return;
  start_read();
  start_write();

  io_context.restart();
  io_context.poll();

  start_read();
  start_write();
}

void UDP_Server::start_read() {
  if (read_pending || !bound || !socket || !socket->is_open())
    return;
  if (recv_storage.empty())
    recv_storage.assign(static_cast<size_t>(read_chunk_size), 0);

  read_pending = true;
  socket->async_receive_from(
      asio::buffer(recv_storage.data(), recv_storage.size()), recv_endpoint,
      [this](const asio::error_code &ec, std::size_t n) {
        read_pending = false;
        if (!bound)
          return;
        if (ec == asio::error::operation_aborted)
          return;
        if (ec)
          return;

        last_peer = endpoint_to_sockaddr(recv_endpoint);
        has_last_peer = true;
        if (std::find(clients.begin(), clients.end(), last_peer) ==
            clients.end()) {
          clients.push_back(last_peer);
        }

        if (recv_buffer.write(recv_storage.data(), n)) {
          recv_peers.push_back(last_peer);
        }
        start_read();
      });
}

void UDP_Server::start_write() {
  if (write_pending || !bound || !socket || !socket->is_open())
    return;
  if (send_endpoints.empty())
    return;
  if (send_storage.empty())
    send_storage.assign(static_cast<size_t>(read_chunk_size), 0);

  std::size_t out_len = send_storage.size();
  if (!send_buffer.peek(send_storage.data(), out_len))
    return;
  send_storage.resize(out_len);
  auto endpoint = send_endpoints.front();

  write_pending = true;
  socket->async_send_to(
      asio::buffer(send_storage.data(), send_storage.size()), endpoint,
      [this](const asio::error_code &ec, std::size_t) {
        write_pending = false;
        if (!bound)
          return;
        if (ec == asio::error::operation_aborted)
          return;
        if (!ec) {
          send_buffer.skip_one();
          if (!send_endpoints.empty())
            send_endpoints.erase(send_endpoints.begin());
        }
        send_storage.assign(static_cast<size_t>(read_chunk_size), 0);
        start_write();
      });
}

void UDP_Server::close() {
  read_pending = false;
  write_pending = false;
  if (socket) {
    asio::error_code ec;
    socket->cancel(ec);
    socket->close(ec);
  }
  clients.clear();
  recv_peers.clear();
  send_endpoints.clear();
  bound = false;
  has_last_peer = false;
  socket_fd = static_cast<Socket_FD>(-1);
  state = Not_Created;
}

std::vector<UDP_Server::Read_Info> UDP_Server::read() {
  std::vector<Read_Info> ret;
  std::string data(static_cast<size_t>(read_chunk_size), '\0');

  while (!recv_peers.empty()) {
    std::size_t out_len = data.size();
    if (!recv_buffer.read(data.data(), out_len))
      break;
    ret.push_back({recv_peers.front(), std::string(data.data(), out_len)});
    recv_peers.erase(recv_peers.begin());
  }
  return ret;
}

void UDP_Server::reply_last_peer(std::string_view data) {
  if (!has_last_peer)
    return;
  send_to(last_peer.ip, static_cast<uint16_t>(last_peer.port), data);
}

void UDP_Server::send_to(std::string_view ip, uint16_t port,
                         std::string_view data) {
  if (data.empty())
    return;
  if (!bound && !bind())
    return;

  asio::error_code ec;
  auto endpoint = asio::ip::udp::endpoint(asio::ip::make_address(ip, ec), port);
  if (ec)
    return;
  if (send_buffer.write(data.data(), data.size())) {
    send_endpoints.push_back(endpoint);
  }
  start_write();
}

void UDP_Server::write_to_all_clients(std::string_view msg) {
  for (const auto &client : clients) {
    send_to(client.ip, static_cast<uint16_t>(client.port), msg);
  }
}

} // namespace Psc::asio_socket