TCP input TLS - Optimize TlsDecoder.

This also brings new api for `DecodeBuffer` and simplification of IpfixDecoder.

Author: BonnyAD9

src/plugins/input/tcp/src/DecodeBuffer.cpp

@@ -17,6 +17,8 @@
 
 #include <ipfixcol2.h> // fds_ipfix_msg_hdr
 
+#include "read_until_n.hpp"
+
 namespace tcp_in {
 
 void DecodeBuffer::read_from(const uint8_t *data, size_t size) {
@@ -36,6 +38,18 @@ void DecodeBuffer::read_from(const uint8_t *data, size_t size) {
     }
 }
 
+void DecodeBuffer::read_from(Reader &reader) {
+    while (!enough_data()) {
+        if (!read_header(reader)) {
+            break;
+        }
+
+        if (!read_body(reader)) {
+            break;
+        }
+    }
+}
+
 void DecodeBuffer::read_from(
     const uint8_t *data,
     size_t buffer_size,
@@ -63,14 +77,15 @@ void DecodeBuffer::signal_eof() {
     m_eof_reached = true;
 }
 
-bool DecodeBuffer::read_header(const uint8_t **data, size_t *size) {
+template<typename... State>
+bool DecodeBuffer::read_header(State &&...state) {
     if (m_decoded_size != 0) {
         // The header is already read, but the message body is incomplete.
         return true;
     }
 
     // Read the header.
-    if (!read_until_n(sizeof(fds_ipfix_msg_hdr), data, size)) {
+    if (!read_until_n(sizeof(fds_ipfix_msg_hdr), std::forward<State>(state)...)) {
         // There is not enough data to read the whole header.
         return false;
     }
@@ -86,9 +101,10 @@ bool DecodeBuffer::read_header(const uint8_t **data, size_t *size) {
     return true;
 }
 
-bool DecodeBuffer::read_body(const uint8_t **data, size_t *size) {
+template<typename... State>
+bool DecodeBuffer::read_body(State &&...state) {
     // Read the body
-    if (!read_until_n(m_decoded_size, data, size)) {
+    if (!read_until_n(m_decoded_size, std::forward<State>(state)...)) {
         // There is not enough data to read the whole body.
         return false;
     }
@@ -107,6 +123,10 @@ bool DecodeBuffer::read_until_n(size_t n, const uint8_t **data, size_t *data_len
     return m_part_decoded.size() == n;
 }
 
+bool DecodeBuffer::read_until_n(size_t n, Reader &reader) {
+    return ::read_until_n(n, reader, m_part_decoded, *this);
+}
+
 size_t DecodeBuffer::read_min(const uint8_t *data, size_t size1, size_t size2) {
     auto size = std::min(size1, size2);
     auto filled = m_part_decoded.size();

src/plugins/input/tcp/src/DecodeBuffer.hpp

@@ -10,11 +10,12 @@
 
 #pragma once
 
-#include <vector>  // std::vector
 #include <cstdint> // uint8_t, UINT16_MAX
 #include <cstddef> // size_t
+#include <vector>  // std::vector
 
 #include "ByteVector.hpp" // ByteVector
+#include "Reader.hpp"
 
 namespace tcp_in {
 
@@ -60,6 +61,12 @@ class DecodeBuffer {
      */
     void read_from(const uint8_t *data, size_t size);
 
+    /**
+     * @brief Read data using generic reader function.
+     * @param reader generic reader function.
+     */
+    void read_from(Reader &reader);
+
     /**
      * @brief Copies IPFIX data from circular buffer.
      *
@@ -104,18 +111,18 @@ class DecodeBuffer {
 private:
     /**
      * @brief Reads the length from ipfix header to `m_decoded_size`.
-     * @param[in,out] data Pointer to data. It is modified by the amount of read data.
-     * @param[in,out] size Size of the data. It is modified by the amount of read data.
+     * @param state state of the reader. This may be anything for which is read_unitl_n implemented.
      * @return `true` if the whole header could be read.
      */
-    bool read_header(const uint8_t **data, size_t *size);
+    template<typename... State>
+    bool read_header(State &&...state);
     /**
      * @brief Reads the body of ipfix message.
-     * @param data Pointer to data. It is modified by the amount of read data.
-     * @param size Size of the data. It is modified by the amount of read data.
+     * @param state state of the reader. This may be anything for which is read_unitl_n implemented.
      * @return `true` if the whole header could be read.
      */
-    bool read_body(const uint8_t **data, size_t *size);
+    template<typename... State>
+    bool read_body(State &&...state);
     /**
      * @brief Reads to `m_part_decoded` until it reaches length of `n`.
      * @param n Target length for `m_part_decoded`.
@@ -124,6 +131,13 @@ class DecodeBuffer {
      * @return `true` if `m_part_decoded` has reached length of `n`.
      */
     bool read_until_n(size_t n, const uint8_t **data, size_t *data_len);
+    /**
+     * @brief Reads to `m_part_decoded` until it reaches length of `n`.
+     * @param n Target length for `m_part_decoded`.
+     * @param reader Generic reader from which to read.
+     * @return `true` if `m_part_decoded` has reached length of `n`.
+     */
+    bool read_until_n(size_t n, Reader &reader);
     /**
      * @brief Adds N elements from `src` to `m_part_decoded` where N is the smaller of the two
      * sizes.

src/plugins/input/tcp/src/IpfixDecoder.cpp

@@ -10,77 +10,16 @@
 
 #include "IpfixDecoder.hpp"
 
-#include <stdexcept> // runtime_error
-#include <string>    // string
-#include <errno.h>   // errno, EWOULDBLOCK, EAGAIN
-#include <stddef.h>  // size_t
-
-#include <unistd.h> // read
-
-#include <ipfixcol2.h> // fds_ipfix_msg_hdr, ipx_strerror
-
 #include "DecodeBuffer.hpp" // DecodeBuffer
-#include "read_until_n.hpp" // read_until_n
+#include "Reader.hpp"
 
 namespace tcp_in {
 
 DecodeBuffer &IpfixDecoder::decode() {
-    while (!m_decoded.enough_data()) {
-        // Read the header of the message.
-        if (!read_header()) {
-            // There is not enough data to read the whole header.
-            break;
-        }
-
-        // Read the body of the message.
-        if (!read_body()) {
-            // There is not enough data to read the whole body of the message.
-            break;
-        }
-    }
-
-    if (m_decoded.is_eof_reached() && m_part_readed.size() != 0) {
-        throw std::runtime_error("Received incomplete message.");
-    }
-
+    auto reader = tcp_reader(m_fd);
+    m_decoded.read_from(reader);
     return m_decoded;
 }
 
-bool IpfixDecoder::read_header() {
-    if (m_msg_size != 0) {
-        // The header is already read, but the message body is incomplete.
-        return true;
-    }
-
-    // Read the header
-    if (!read_until_n(sizeof(fds_ipfix_msg_hdr))) {
-        // There is not enough data to read the whole header.
-        return false;
-    }
-
-    // The header has been read, load the size of the message.
-    auto hdr = reinterpret_cast<fds_ipfix_msg_hdr *>(m_part_readed.data());
-    m_msg_size = ntohs(hdr->length);
-    return true;
-}
-
-bool IpfixDecoder::read_body() {
-    // Read the body.
-    if (!read_until_n(m_msg_size)) {
-        // There is not enough data to read the whole body
-        return false;
-    }
-
-    // Whole message has been read. Add it to the decode buffer.
-    m_decoded.add(std::move(m_part_readed));
-    m_part_readed = ByteVector();
-    m_msg_size = 0;
-    return true;
-}
-
-bool IpfixDecoder::read_until_n(size_t n) {
-    return ::read_until_n(n, m_fd, m_part_readed, m_decoded);
-}
-
 } // namespace tcp_in
 

src/plugins/input/tcp/src/IpfixDecoder.hpp

@@ -15,7 +15,6 @@
 
 #include "Decoder.hpp"      // Decoder
 #include "DecodeBuffer.hpp" // DecodeBuffer
-#include "ByteVector.hpp"   // ByteVector
 
 namespace tcp_in {
 
@@ -29,7 +28,7 @@ class IpfixDecoder : public Decoder {
      * @brief Creates ipfix decoder.
      * @param fd TCP connection file descriptor.
      */
-    IpfixDecoder(int fd) : m_fd(fd), m_decoded(), m_part_readed(), m_msg_size(0) {}
+    IpfixDecoder(int fd) : m_fd(fd), m_decoded() {}
 
     virtual DecodeBuffer &decode() override;
 
@@ -38,19 +37,8 @@ class IpfixDecoder : public Decoder {
     }
 
 private:
-    /** returns true if there was enough data to read the header or if the header is already read */
-    bool read_header();
-    /** returns true if there was enough data to read the body */
-    bool read_body();
-    /** returns true if there was enough data to read to the given amount */
-    bool read_until_n(size_t n);
-
     int m_fd;
     DecodeBuffer m_decoded;
-
-    ByteVector m_part_readed;
-    /** Final size of the whole IPFIX message. When 0 IPIFX header is not fully read. */
-    size_t m_msg_size;
 };
 
 } // namespace tcp_in

src/plugins/input/tcp/src/Lz4Decoder.cpp

@@ -177,7 +177,8 @@ void Lz4Decoder::decompress() {
 }
 
 bool Lz4Decoder::read_until_n(size_t n) {
-    return ::read_until_n(n, m_fd, m_compressed, m_decoded);
+    auto reader = tcp_reader(m_fd);
+    return ::read_until_n(n, reader, m_compressed, m_decoded);
 }
 
 void Lz4Decoder::reset_stream(size_t buffer_size) {

src/plugins/input/tcp/src/Reader.hpp

@@ -0,0 +1,52 @@
+/**
+ * \file
+ * \author Jakub Antonín Štigler <[email protected]>
+ * \brief Generic reader function. (header file)
+ * \date 2025
+ *
+ * Copyright: (C) 2023 CESNET, z.s.p.o.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#pragma once
+
+#include <cerrno>
+#include <cstdint>
+#include <functional>
+#include <stdexcept>
+
+#include <unistd.h>
+#include <sys/socket.h>
+
+#include <ipfixcol2.h>
+
+/** Describes result of read operation. */
+enum class ReadResult {
+    /** Successufully read some data. */
+    READ,
+    /** Non blocking socket needs to wait for more data. */
+    WAIT,
+    /** Connection has ended. */
+    END,
+};
+
+using Reader = std::function<ReadResult(std::uint8_t *data, std::size_t &length)>;
+
+static inline Reader tcp_reader(int fd) {
+    return [=](std::uint8_t *data, std::size_t &length) -> ReadResult {
+        int res = recv(fd, data, length, 0);
+        if (res == -1) {
+            length = 0;
+            int err = errno;
+            if (err == EWOULDBLOCK || err == EAGAIN) {
+                return ReadResult::WAIT;
+            }
+            const char *err_str;
+            ipx_strerror(err, err_str);
+            throw std::runtime_error("Failed to read from descruptor " + std::string(err_str));
+        }
+
+        length = std::size_t(res);
+        return length == 0 ? ReadResult::END : ReadResult::READ;
+    };
+}

src/plugins/input/tcp/src/read_until_n.hpp

@@ -10,59 +10,44 @@
 
 #pragma once
 
-#include <cstddef>   // size_t
 #include <cerrno>    // errno, EWOULDBLOCK, EAGAIN
-#include <string>    // std::string
+#include <cstddef>   // size_t
 #include <stdexcept> // std::runtime_error
-
-#include <sys/socket.h> // recv
+#include <string>    // std::string
 
 #include <ipfixcol2.h> // ipx_strerror
 
+#include "Reader.hpp"
+
 /**
  * @brief Reads to vector from file descriptor until the vector has the desired number of bytes
  *
  * @tparam Vec vector like object
  * @tparam EofSig `EofSig::signal_eof()` is called when eof is reached
  * @param n number of bytes that should be in the vector
- * @param fd file descriptor to read from
+ * @param reader generic read function
  * @param result where to read the bytes to
  * @param eofSignaler has method to signal eof
  * @return true if the vector has at least `n` bytes after the call
  * @throws when fails to read from the file descriptor
  */
 template<typename Vec, typename EofSig>
-bool read_until_n(size_t n, int fd, Vec &result, EofSig &eofSignaler) {
+bool read_until_n(size_t n, Reader &reader, Vec &result, EofSig &eofSignaler) {
     auto filled = result.size();
     if (filled >= n) {
         return true;
     }
 
-    auto remaining = n - filled;
     result.resize(n);
 
-    int res = recv(fd, result.data() + filled, remaining, 0);
-    if (res == -1) {
-        result.resize(filled);
-        int err = errno;
-        if (err == EWOULDBLOCK || err == EAGAIN) {
-            return false;
-        }
-
-        const char *err_str;
-        ipx_strerror(err, err_str);
-        throw std::runtime_error("Failed to read from descriptor: " + std::string(err_str));
-    }
+    auto remaining = n - filled;
+    auto res = reader(result.data() + filled, remaining);
 
-    result.resize(filled + res);
+    result.resize(filled + remaining);
 
-    if (res == 0) {
+    if (res == ReadResult::END) {
         eofSignaler.signal_eof();
     }
 
-    if (static_cast<size_t>(res) != remaining) {
-        return false;
-    }
-
-    return true;
+    return result.size() == n;
 }