control: Limit memory use by control connections

Signed-off-by: Mobin "Hojjat" Aydinfar <mobin@mobintestserver.ir>

src/control.cc

@@ -1135,7 +1135,8 @@ void control_conn_t::service_event(service_record *service, service_event_t even
 
 bool control_conn_t::queue_packet(const char *pkt, unsigned size) noexcept
 {
-    int in_flag = bad_conn_close ? 0 : IN_EVENTS;
+    // Close connection on bad con or temporarily Limit input if necessary
+    int in_flag = bad_conn_close || (outbuf_size >= 1000) ? 0 : IN_EVENTS;
     bool was_empty = outbuf.empty();
 
     // If the queue is empty, we can try to write the packet out now rather than queueing it.
@@ -1166,6 +1167,7 @@ bool control_conn_t::queue_packet(const char *pkt, unsigned size) noexcept
     // Create a vector out of the (remaining part of the) packet:
     try {
         outbuf.emplace_back(pkt, pkt + size);
+        outbuf_size += size;
         iob.set_watches(in_flag | OUT_EVENTS);
         return true;
     }
@@ -1190,7 +1192,8 @@ bool control_conn_t::queue_packet(const char *pkt, unsigned size) noexcept
 // make them extraordinary difficult to combine into a single method.
 bool control_conn_t::queue_packet(std::vector<char> &&pkt) noexcept
 {
-    int in_flag = bad_conn_close ? 0 : IN_EVENTS;
+    // Close connection on bad con or temporarily Limit input if necessary
+    int in_flag = bad_conn_close || (outbuf_size >= 1000) ? 0 : IN_EVENTS;
     bool was_empty = outbuf.empty();
 
     if (was_empty) {
@@ -1219,6 +1222,7 @@ bool control_conn_t::queue_packet(std::vector<char> &&pkt) noexcept
 
     try {
         outbuf.emplace_back(std::move(pkt));
+        outbuf_size += pkt.size();
         iob.set_watches(in_flag | OUT_EVENTS);
         return true;
     }
@@ -1275,8 +1279,10 @@ bool control_conn_t::data_ready() noexcept
         iob.set_watches(OUT_EVENTS);
     }
     else {
+        // Temporarily Limit input if necessary
+        int in_flags = (outbuf_size < 1000) ? IN_EVENTS : 0;
         int out_flags = (bad_conn_close || !outbuf.empty()) ? OUT_EVENTS : 0;
-        iob.set_watches(IN_EVENTS | out_flags);
+        iob.set_watches(in_flags | out_flags);
     }
 
     return false;
@@ -1314,6 +1320,7 @@ bool control_conn_t::send_data() noexcept
     outpkt_index += written;
     if (outpkt_index == pkt.size()) {
         // We've finished this packet, move on to the next:
+        outbuf_size -= pkt.size();
         outbuf.pop_front();
         outpkt_index = 0;
         if (oom_close) {
@@ -1324,7 +1331,9 @@ bool control_conn_t::send_data() noexcept
             if (bad_conn_close) {
                 return true;
             }
-            iob.set_watches(IN_EVENTS);
+            // Temporarily Limit input if necessary
+            int in_flag = (outbuf_size < 1000) ? IN_EVENTS : 0;
+            iob.set_watches(in_flag);
         }
     }
 

src/includes/control.h

@@ -111,6 +111,9 @@ class control_conn_t : private service_listener
 
     // Buffer for outgoing packets. Each outgoing back is represented as a vector<char>.
     list<vector<char>> outbuf;
+    // Output buffer size. functions will increase this value when pushing something into
+    // outbuf and lowering this value when something get removed from outbuf.
+    unsigned outbuf_size = 0;
     // Current index within the first outgoing packet (all previous bytes have been sent).
     unsigned outpkt_index = 0;
 
@@ -181,6 +184,9 @@ class control_conn_t : private service_listener
     bool data_ready() noexcept;
 
     bool send_data() noexcept;
+
+    // Check the output buffer size and limit inputs if necessary
+    bool input_memlimit() noexcept;
 
     // Check if any dependents will be affected by stopping a service, generate a response packet if so.
     // had_dependents will be set true if the service should not be stopped, false otherwise.