added handshake for dos protection

repertory/librepertory/include/comm/packet/packet_client.hpp

@@ -67,6 +67,8 @@ private:
 
   [[nodiscard]] auto get_client() -> std::shared_ptr<client>;
 
+  [[nodiscard]] auto handshake(client &cli) const -> bool;
+
   void put_client(std::shared_ptr<client> &cli);
 
   [[nodiscard]] auto read_packet(client &cli, packet &response) const

repertory/librepertory/include/comm/packet/packet_server.hpp

@@ -77,6 +77,8 @@ private:
 private:
   void add_client(connection &conn, const std::string &client_id);
 
+  [[nodiscard]] auto handshake(std::shared_ptr<connection> conn) const -> bool;
+
   void initialize(const uint16_t &port, uint8_t pool_size);
 
   void listen_for_connection(tcp::acceptor &acceptor);

repertory/librepertory/src/comm/packet/packet_client.cpp

@@ -97,11 +97,11 @@ struct non_blocking_guard final {
   return false;
 }
 
-template <class op_t, class cancel_t>
+template <class op_t>
-void run_with_deadline(boost::asio::io_context &io_ctx, op_t &&operation,
+void run_with_deadline(boost::asio::io_context &io_ctx,
-                       cancel_t &&cancel_op, std::chrono::milliseconds deadline,
+                       boost::asio::ip::tcp::socket &sock, op_t &&operation,
-                       std::string_view timeout_event_tag,
+                       std::chrono::milliseconds deadline,
-                       std::string_view op_name,
+                       std::string_view event_name,
                        std::string_view function_name) {
   deadline = std::max(deadline, std::chrono::milliseconds{250});
 
@@ -110,10 +110,10 @@ void run_with_deadline(boost::asio::io_context &io_ctx, op_t &&operation,
 
   boost::asio::steady_timer timer{io_ctx};
   timer.expires_after(deadline);
-  timer.async_wait([&cancel_op, &done, &timed_out](auto &&err_) {
+  timer.async_wait([&done, &sock, &timed_out](auto &&err_) {
     if (not err_ && not done) {
       timed_out = true;
-      cancel_op();
+      sock.cancel();
     }
   });
 
@@ -131,12 +131,13 @@ void run_with_deadline(boost::asio::io_context &io_ctx, op_t &&operation,
 
   if (timed_out) {
     repertory::event_system::instance().raise<repertory::packet_client_timeout>(
-        std::string(timeout_event_tag), std::string(function_name));
+        std::string(event_name), std::string(function_name));
-    throw std::runtime_error(std::string(op_name) + " timed-out");
+    throw std::runtime_error(std::string(event_name) + " timed-out");
   }
 
   if (err) {
-    throw std::runtime_error(std::string(op_name) + " failed|" + err.message());
+    throw std::runtime_error(std::string(event_name) + " failed|err|" +
+                             err.message());
   }
 }
 
@@ -147,13 +148,12 @@ void connect_with_deadline(boost::asio::io_context &io_ctx,
   REPERTORY_USES_FUNCTION_NAME();
 
   run_with_deadline(
-      io_ctx,
+      io_ctx, sock,
       [&sock, &endpoints](auto &&handler) {
         boost::asio::async_connect(
             sock, endpoints, [handler](auto &&err, auto &&) { handler(err); });
       },
-      [&sock]() { sock.cancel(); }, deadline, "connect", "connect",
+      deadline, "connect", function_name);
-      function_name);
 }
 
 void read_exact_with_deadline(boost::asio::io_context &io_ctx,
@@ -171,7 +171,7 @@ void read_exact_with_deadline(boost::asio::io_context &io_ctx,
     std::size_t bytes_read = 0U;
 
     run_with_deadline(
-        io_ctx,
+        io_ctx, sock,
         [&](auto &&handler) {
           sock.async_read_some(
               boost::asio::buffer(base + offset, total - offset),
@@ -180,8 +180,7 @@ void read_exact_with_deadline(boost::asio::io_context &io_ctx,
                 handler(err);
               });
         },
-        [&sock]() { sock.cancel(); }, deadline, "response", "read",
+        deadline, "read", function_name);
-        function_name);
 
     if (bytes_read == 0U) {
       throw std::runtime_error("0 bytes read");
@@ -205,7 +204,7 @@ void write_all_with_deadline(boost::asio::io_context &io_ctx,
     std::size_t bytes_written = 0U;
 
     run_with_deadline(
-        io_ctx,
+        io_ctx, sock,
         [&](auto &&handler) {
           sock.async_write_some(
               boost::asio::buffer(base + offset, total - offset),
@@ -214,8 +213,7 @@ void write_all_with_deadline(boost::asio::io_context &io_ctx,
                 handler(err);
               });
         },
-        [&sock]() { sock.cancel(); }, deadline, "request", "write",
+        deadline, "write", function_name);
-        function_name);
 
     if (bytes_written == 0U) {
       throw std::runtime_error("0 bytes written");
@@ -269,10 +267,15 @@ void packet_client::connect(client &cli) {
     cli.socket.set_option(boost::asio::socket_base::linger(false, 0));
     cli.socket.set_option(boost::asio::socket_base::keep_alive(true));
 
+    if (not handshake(cli)) {
+      close(cli);
+      return;
+    }
+
     packet response;
     auto res = read_packet(cli, response);
     if (res != 0) {
-      throw std::runtime_error(std::to_string(res));
+      throw std::runtime_error(fmt::format("read packet failed|err|{}", res));
     }
   } catch (...) {
     close(cli);
@@ -308,6 +311,52 @@ auto packet_client::get_client() -> std::shared_ptr<packet_client::client> {
   return nullptr;
 }
 
+auto packet_client::handshake(client &cli) const -> bool {
+  REPERTORY_USES_FUNCTION_NAME();
+
+  try {
+    data_buffer buffer;
+    {
+      packet tmp;
+      tmp.encode_top(cli.nonce);
+      tmp.to_buffer(buffer);
+    }
+    auto to_read{buffer.size()};
+
+    std::uint32_t total_read{};
+    while ((total_read < to_read) && cli.socket.is_open()) {
+      auto bytes_read = boost::asio::read(
+          cli.socket,
+          boost::asio::buffer(&buffer[total_read], buffer.size() - total_read));
+      if (bytes_read <= 0) {
+        throw std::runtime_error("0 bytes read");
+      }
+
+      total_read += static_cast<std::uint32_t>(bytes_read);
+    }
+
+    if (total_read == to_read) {
+      packet response(buffer);
+      response.encrypt(cfg_.encryption_token);
+      response.to_buffer(buffer);
+
+      auto written = boost::asio::write(
+          cli.socket, boost::asio::buffer(boost::asio::buffer(buffer)));
+      if (written == to_read) {
+        return true;
+      }
+
+      throw std::runtime_error("failed to send handshake");
+    }
+
+    throw std::runtime_error("failed to read handshake");
+  } catch (const std::exception &e) {
+    repertory::utils::error::raise_error(function_name, e, "handlshake failed");
+  }
+
+  return false;
+}
+
 void packet_client::put_client(std::shared_ptr<client> &cli) {
   if (not cli || not is_socket_still_alive(cli->socket)) {
     return;

repertory/librepertory/src/comm/packet/packet_server.cpp

@@ -79,6 +79,67 @@ void packet_server::add_client(connection &conn, const std::string &client_id) {
   }
 }
 
+auto packet_server::handshake(std::shared_ptr<connection> conn) const -> bool {
+  REPERTORY_USES_FUNCTION_NAME();
+
+  try {
+    conn->generate_nonce();
+
+    data_buffer buffer;
+    packet request;
+    request.encode_top(conn->nonce);
+    request.to_buffer(buffer);
+    auto to_read{buffer.size()};
+
+    auto written = boost::asio::write(
+        conn->socket, boost::asio::buffer(boost::asio::buffer(buffer)));
+    if (written == to_read) {
+      conn->buffer.resize(to_read);
+
+      std::uint32_t total_read{};
+      while ((total_read < to_read) && conn->socket.is_open()) {
+        auto bytes_read = boost::asio::read(
+            conn->socket,
+            boost::asio::buffer(&conn->buffer[total_read],
+                                conn->buffer.size() - total_read));
+        if (bytes_read <= 0) {
+          throw std::runtime_error("0 bytes read");
+        }
+
+        total_read += static_cast<std::uint32_t>(bytes_read);
+      }
+
+      if (total_read == to_read) {
+        packet response(conn->buffer);
+        if (response.decrypt(encryption_token_) == 0) {
+          std::string nonce;
+          if (response.decode(nonce) == 0) {
+            if (nonce == conn->nonce) {
+              conn->generate_nonce();
+              return true;
+            }
+
+            throw std::runtime_error("invalid nonce");
+          }
+
+          throw std::runtime_error("invalid nonce");
+        }
+
+        throw std::runtime_error("decryption failed");
+      }
+
+      throw std::runtime_error("invalid handshake");
+    }
+
+    throw std::runtime_error("failed to send handshake");
+  } catch (const std::exception &e) {
+    repertory::utils::error::raise_error(function_name, e, "handlshake failed");
+  }
+
+  conn->socket.close();
+  return false;
+}
+
 void packet_server::initialize(const uint16_t &port, uint8_t pool_size) {
   REPERTORY_USES_FUNCTION_NAME();
 
@@ -123,15 +184,19 @@ void packet_server::on_accept(std::shared_ptr<connection> conn,
   if (err) {
     utils::error::raise_error(function_name, err.message());
     std::this_thread::sleep_for(1s);
-  } else {
+    return;
+  }
+
   conn->socket.set_option(boost::asio::ip::tcp::no_delay(true));
   conn->socket.set_option(boost::asio::socket_base::linger(false, 0));
 
-    conn->generate_nonce();
+  if (not handshake(conn)) {
+    conn->socket.close();
+    return;
+  }
 
   packet response;
   send_response(conn, 0, response);
-  }
 }
 
 void packet_server::read_header(std::shared_ptr<connection> conn) {
@@ -145,12 +210,12 @@ void packet_server::read_header(std::shared_ptr<connection> conn) {
         if (err) {
           remove_client(*conn);
           repertory::utils::error::raise_error(function_name, err.message());
-        } else {
+          return;
-          auto to_read =
+        }
-              *reinterpret_cast<std::uint32_t *>(conn->buffer.data());
+
+        auto to_read = *reinterpret_cast<std::uint32_t *>(conn->buffer.data());
         boost::endian::big_to_native_inplace(to_read);
         read_packet(conn, to_read);
-        }
       });
 }