diff --git a/include/TCPClient.h b/include/TCPClient.h index 54754dc..f90ed95 100644 --- a/include/TCPClient.h +++ b/include/TCPClient.h @@ -31,8 +31,8 @@ class TCPClient final : public ICommunicationClient { public: TCPClient(std::string ip, const std::shared_ptr>>> &rx_queue) - : port{3001}, m_ip{std::move(ip)}, m_stop_flag(false), - m_thread(std::thread(&TCPClient::rx_thread, this)), m_rx_queue(rx_queue) { + : port{3001}, m_ip{std::move(ip)}, m_stop_flag(false), m_rx_queue(rx_queue) { + m_thread = std::thread(&TCPClient::rx_thread, this); } ~TCPClient() override; int init() override; diff --git a/include/UDPClient.h b/include/UDPClient.h index 4e9fdd7..381c6ba 100644 --- a/include/UDPClient.h +++ b/include/UDPClient.h @@ -31,8 +31,8 @@ class UDPClient final : public ICommunicationClient { public: UDPClient(std::string ip, const std::shared_ptr>>> &rx_queue) - : m_stop_flag(false), m_thread(std::thread(&UDPClient::rx_thread, this)), - m_remote_ip(std::move(ip)), m_rx_queue(rx_queue) { + : m_stop_flag(false), m_remote_ip(std::move(ip)), m_rx_queue(rx_queue) { + m_thread = std::thread(&UDPClient::rx_thread, this); } ~UDPClient() override; int init() override; diff --git a/include/librpc.h b/include/librpc.h index 5e3f3ef..0d9b5c5 100644 --- a/include/librpc.h +++ b/include/librpc.h @@ -24,8 +24,7 @@ struct SizeAndSource { class MessagingInterface { public: MessagingInterface() - : m_stop_flag(false), m_rx_thread(std::thread(&MessagingInterface::handle_recv, this)), - m_fn_rx_thread(std::thread(&MessagingInterface::handle_fn_recv, this)), + : m_stop_flag(false), m_rx_queue(std::make_shared>>>( RX_QUEUE_SIZE)) { #ifdef _WIN32 @@ -34,6 +33,8 @@ class MessagingInterface { #endif // Initialization must be after call to WSAStartup m_discovery_service = std::make_unique(); + m_rx_thread = std::thread(&MessagingInterface::handle_recv, this); + m_fn_rx_thread = std::thread(&MessagingInterface::handle_fn_recv, this); } ~MessagingInterface(); diff --git a/src/mDNSDiscoveryService.cpp b/src/mDNSDiscoveryService.cpp index 57dcac2..da8cef8 100644 --- a/src/mDNSDiscoveryService.cpp +++ b/src/mDNSDiscoveryService.cpp @@ -217,6 +217,9 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf const int size) { int ptr = 0; mDNSRobotModule response{}; + const auto remaining = [&](const int needed) { + return needed >= 0 && ptr >= 0 && ptr <= size && needed <= size - ptr; + }; // Header if (size < sizeof(query_header)) { @@ -231,7 +234,7 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf // Questions for (int i = 0; i < h->num_questions; i++) { - if (ptr > size) { + if (ptr >= size) { return std::nullopt; } @@ -241,13 +244,16 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf return std::nullopt; } ptr = new_ptr; + if (!remaining(static_cast(sizeof(query_footer)))) { + return std::nullopt; + } ptr += sizeof(query_footer); } // Answers and authority (we do not care about authority). bool robot_module = false; for (int i = 0; i < h->num_answers + h->num_authority + h->num_additional; i++) { - if (ptr > size) { + if (ptr >= size) { return std::nullopt; } // We assume that the boards mdns does not send any questions asking for @@ -263,6 +269,9 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf robot_module |= name.find("_robotcontrol") != std::string::npos; response.hostname = name; + if (!remaining(static_cast(sizeof(answer)))) { + return std::nullopt; + } const auto a = reinterpret_cast(buffer + ptr); a->type = ntohs(a->type); a->answer_class = ntohs(a->answer_class); @@ -270,6 +279,10 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf a->data_length = ntohs(a->data_length); ptr += sizeof(answer); + if (!remaining(a->data_length)) { + return std::nullopt; + } + // A-Record if (a->type == 1 && robot_module) { std::vector data; @@ -290,7 +303,13 @@ std::optional mDNSDiscoveryService::parse_response(uint8_t *buf if (a->type == 16 && robot_module) { int inner_ptr = ptr; while (inner_ptr < a->data_length + ptr) { + if (inner_ptr >= size) { + return std::nullopt; + } const int len = buffer[inner_ptr++]; + if (len < 0 || inner_ptr > size - len || inner_ptr > ptr + a->data_length - len) { + return std::nullopt; + } std::string s(reinterpret_cast(buffer + inner_ptr), len); inner_ptr += len; @@ -347,7 +366,7 @@ std::tuple mDNSDiscoveryService::read_mdns_name(const uint8_t if (buffer[ptr] >= 0xC0) { // compressed ptr++; - if (buffer[ptr] < 0 || buffer[ptr] > ptr) { + if (ptr >= size || buffer[ptr] > ptr) { return {"", -1}; } const auto [name, l] = read_mdns_name(buffer, size, buffer[ptr]);