control: Limit memory use by control connections.

a new call named `DINIT_CP_CAN` has been added to accept or reject
new requests. Dinit will reply with DINIT_RP_IDLE or DINIT_RP_BUSY.

Signed-off-by: Mobin "Hojjat" Aydinfar <[email protected]>

src/control.cc

@@ -12,7 +12,7 @@
 
 // Control protocol versions:
 // 1 - dinit 0.16 and prior
-// 2 - dinit 0.17 (adds DINIT_CP_SETTRIGGER, DINIT_CP_CATLOG)
+// 2 - dinit 0.17 (adds DINIT_CP_SETTRIGGER, DINIT_CP_CATLOG, DINIT_CP_CAN)
 
 namespace {
     constexpr auto OUT_EVENTS = dasynq::OUT_EVENTS;
@@ -40,6 +40,21 @@ bool control_conn_t::process_packet()
     // shouldn't touch instance members after that point.
 
     int pktType = rbuf[0];
+    if (pktType == DINIT_CP_CAN) {
+        // Responds with
+        // DINIT_RP_BUSY or DINIT_RP_IDLE
+        if (sizeof(outbuf) < 1000) {
+            char replyBuf[] = { DINIT_RP_IDLE };
+            if (! queue_packet(replyBuf, sizeof(replyBuf))) return false;
+            rbuf.consume(1);
+            return true;
+        } else {
+            char replyBuf[] = { DINIT_RP_BUSY };
+            if (! queue_packet(replyBuf, sizeof(replyBuf))) return false;
+            rbuf.consume(1);
+            return true;
+        }
+    }
     if (pktType == DINIT_CP_QUERYVERSION) {
         // Responds with:
         // DINIT_RP_CVERSION, (2 byte) minimum compatible version, (2 byte) actual version

src/dinitctl.cc

@@ -410,6 +410,7 @@ int dinitctl_main(int argc, char **argv)
         // Start by querying protocol version:
         cpbuffer_t rbuffer;
         uint16_t daemon_protocol_ver = check_protocol_version(min_cp_version, max_cp_version, rbuffer, socknum);
+        if (daemon_protocol_ver >= 2) check_dinit_availability(rbuffer, socknum);
 
         if (command == command_t::UNPIN_SERVICE) {
             return unpin_service(socknum, rbuffer, service_name, verbose);

src/includes/control-cmds.h

@@ -60,6 +60,8 @@ constexpr static int DINIT_CP_SETTRIGGER = 19;
 // Retrieve buffered output
 constexpr static int DINIT_CP_CATLOG = 20;
 
+// Is dinit daemon busy?
+constexpr static int DINIT_CP_CAN = 21;
 
 // Replies:
 
@@ -125,6 +127,10 @@ constexpr static int DINIT_RP_SERVICE_LOAD_ERR = 72;
 // Service log:
 constexpr static int DINIT_RP_SERVICE_LOG = 73;
 
+// Daemon can or cannot not accept more connections yet:
+constexpr static int DINIT_RP_IDLE = 74;
+constexpr static int DINIT_RP_BUSY = 75;
+
 
 // Information (out-of-band):
 

src/includes/dinit-client.h

@@ -38,6 +38,12 @@ class cp_old_server_exception
     // no body
 };
 
+class cp_busy_exception
+{
+    public:
+    cp_busy_exception() { }
+};
+
 class general_error
 {
     int err; // related errno (or 0)
@@ -229,6 +235,21 @@ template <typename Buf> inline void write_all_x(int fd, const Buf &b)
     write_all_x(fd, b.data(), b.size());
 }
 
+// Check the dinit daemon can responde to us.
+// doesn't do anything if dinit is idle (DINIT_RP_IDLE) or
+// throw a cp_busy_exception if dinit is busy (DINIT_RP_BUSY) or
+// throw a cp_read_exception on protocol error
+inline void check_dinit_availability(cpbuffer_t &rbuffer, int fd) {
+    constexpr int bufsize = 1;
+    char buf[bufsize] = { DINIT_CP_CAN };
+    write_all_x(fd, buf, bufsize);
+
+    wait_for_reply(rbuffer, fd);
+    if (rbuffer[0] == DINIT_RP_BUSY) throw cp_busy_exception{};
+    if (rbuffer[0] != DINIT_RP_IDLE) throw cp_read_exception{0};
+    rbuffer.consume(bufsize);
+}
+
 // Check the protocol version is compatible with the client.
 //   minversion - minimum protocol version that client can speak
 //   version - maximum protocol version that client can speak

src/tests/cptests/cptests.cc

@@ -25,6 +25,48 @@ class control_conn_t_test
 // Size of status buffer, as returned in several packet types
 constexpr static int STATUS_BUFFER_SIZE = 6 + ((sizeof(pid_t) > sizeof(int)) ? sizeof(pid_t) : sizeof(int));
 
+void cptest_can()
+{
+    service_set sset;
+    int fd = bp_sys::allocfd();
+    auto *cc = new control_conn_t(event_loop, &sset, fd);
+
+    bp_sys::supply_read_data(fd, { DINIT_CP_CAN });
+
+    event_loop.regd_bidi_watchers[fd]->read_ready(event_loop, fd);
+
+    // We expect a DINIT_RP_IDLE:
+    std::vector<char> wdata;
+    bp_sys::extract_written_data(fd, wdata);
+
+    assert(wdata.size() == 1);
+    assert(wdata[0] == DINIT_RP_IDLE);
+
+    delete cc;
+}
+/*
+    TODO: How fill entire output buffer for testing?
+
+void cptest_can2()
+{
+    service_set sset;
+    int fd = bp_sys::allocfd();
+    auto *cc = new control_conn_t(event_loop, &sset, fd);
+
+    bp_sys::supply_read_data(fd, { DINIT_CP_CAN });
+
+    event_loop.regd_bidi_watchers[fd]->read_ready(event_loop, fd);
+
+    // We expect a DINIT_RP_BUSY:
+    std::vector<char> wdata;
+    bp_sys::extract_written_data(fd, wdata);
+
+    assert(wdata.size() == 1);
+    assert(wdata[0] == DINIT_RP_BUSY);
+
+    delete cc;
+}
+*/
 void cptest_queryver()
 {
     service_set sset;
@@ -959,6 +1001,8 @@ void cptest_servicestatus()
 
 int main(int argc, char **argv)
 {
+    RUN_TEST(cptest_can, "                ");
+    //RUN_TEST(cptest_can2, "               "); FIXME
     RUN_TEST(cptest_queryver, "           ");
     RUN_TEST(cptest_listservices, "       ");
     RUN_TEST(cptest_findservice1, "       ");