diff options
| -rw-r--r-- | include/salticidae/conn.h | 17 | ||||
| -rw-r--r-- | include/salticidae/network.h | 15 | ||||
| -rw-r--r-- | include/salticidae/util.h | 13 | ||||
| -rw-r--r-- | src/util.cpp | 3 |
4 files changed, 40 insertions, 8 deletions
diff --git a/include/salticidae/conn.h b/include/salticidae/conn.h index 2f72376..4e2f0f9 100644 --- a/include/salticidae/conn.h +++ b/include/salticidae/conn.h @@ -358,10 +358,13 @@ class ConnPool { public: class Config { - friend ConnPool; + friend class ConnPool; + template<typename OpcodeType> friend class MsgNetwork; int _max_listen_backlog; double _conn_server_timeout; size_t _seg_buff_size; + size_t _max_msg_size; + size_t _max_msg_queue_size; size_t _max_recv_buff_size; size_t _nworker; size_t _queue_capacity; @@ -378,6 +381,8 @@ class ConnPool { _max_listen_backlog(10), _conn_server_timeout(2), _seg_buff_size(4096), + _max_msg_size(1024), + _max_msg_queue_size(65536), _max_recv_buff_size(4096), _nworker(1), _queue_capacity(0), @@ -404,6 +409,16 @@ class ConnPool { return *this; } + Config &max_msg_size(size_t x) { + _max_msg_size = x; + return *this; + } + + Config &max_msg_queue_size(size_t x) { + _max_msg_queue_size = x; + return *this; + } + Config &max_recv_buff_size(size_t x) { _max_recv_buff_size = x; return *this; diff --git a/include/salticidae/network.h b/include/salticidae/network.h index 48b555f..18f3a42 100644 --- a/include/salticidae/network.h +++ b/include/salticidae/network.h @@ -120,6 +120,8 @@ class MsgNetwork: public ConnPool { #endif private: + const size_t max_msg_size; + const size_t max_msg_queue_size; std::unordered_map< typename Msg::opcode_t, std::function<void(const Msg &msg, const conn_t &)>> handler_map; @@ -170,8 +172,10 @@ class MsgNetwork: public ConnPool { virtual ~MsgNetwork() { stop(); } MsgNetwork(const EventContext &ec, const Config &config): - ConnPool(ec, config) { - incoming_msgs.set_capacity(65536); + ConnPool(ec, config), + max_msg_size(config._max_msg_size), + max_msg_queue_size(config._max_msg_queue_size) { + incoming_msgs.set_capacity(max_msg_queue_size); incoming_msgs.reg_handler(ec, [this, burst_size=config._burst_size](queue_t &q) { std::pair<Msg, conn_t> item; size_t cnt = 0; @@ -560,6 +564,13 @@ void MsgNetwork<OpcodeType>::on_read(const ConnPool::conn_t &_conn) { if (recv_buffer.size() < Msg::header_size) break; /* new header available */ msg = Msg(recv_buffer.pop(Msg::header_size)); + if (msg.get_length() > max_msg_size) + { + SALTICIDAE_LOG_WARN( + "oversized message from %s, terminating the connection", + std::string(*conn).c_str()); + throw MsgNetworkError(SALTI_ERROR_CONN_OVERSIZED_MSG); + } msg_state = Conn::PAYLOAD; } if (msg_state == Conn::PAYLOAD) diff --git a/include/salticidae/util.h b/include/salticidae/util.h index 3196042..b2ddd8e 100644 --- a/include/salticidae/util.h +++ b/include/salticidae/util.h @@ -106,7 +106,8 @@ enum SalticidaeErrorCode { SALTI_ERROR_RAND_SOURCE, SALTI_ERROR_CONN_NOT_READY, SALTI_ERROR_NOT_AVAIL, - SALTI_ERROR_UNKNOWN + SALTI_ERROR_UNKNOWN, + SALTI_ERROR_CONN_OVERSIZED_MSG }; extern const char *SALTICIDAE_ERROR_STRINGS[]; @@ -148,12 +149,16 @@ struct ConnPoolError: public SalticidaeError { using SalticidaeError::SalticidaeError; }; -class PeerNetworkError: public ConnPoolError { +class MsgNetworkError: public ConnPoolError { using ConnPoolError::ConnPoolError; }; -class ClientNetworkError: public ConnPoolError { - using ConnPoolError::ConnPoolError; +class PeerNetworkError: public MsgNetworkError { + using MsgNetworkError::MsgNetworkError; +}; + +class ClientNetworkError: public MsgNetworkError { + using MsgNetworkError::MsgNetworkError; }; extern const char *TTY_COLOR_RED; diff --git a/src/util.cpp b/src/util.cpp index ce98b22..8ca01aa 100644 --- a/src/util.cpp +++ b/src/util.cpp @@ -63,7 +63,8 @@ const char *SALTICIDAE_ERROR_STRINGS[] = { "rand source is not available, try again", "connection is not ready", "operation not available", - "unknown error" + "unknown error", + "oversized message", }; const char *TTY_COLOR_RED = "\x1b[31m"; |