src/control.cc

   1 #include <algorithm>
   2 #include <unordered_set>
   3 #include <climits>
   4
   5 #include "control.h"
   6 #include "service.h"
   7
   8 // Server-side control protocol implementation. This implements the functionality that allows
   9 // clients (such as dinitctl) to query service state and issue commands to control services.
  10
  11 namespace {
  12     constexpr auto OUT_EVENTS = dasynq::OUT_EVENTS;
  13     constexpr auto IN_EVENTS = dasynq::IN_EVENTS;
  14
  15     // Control protocol minimum compatible version and current version:
  16     constexpr uint16_t min_compat_version = 1;
  17     constexpr uint16_t cp_version = 1;
  18
  19     // check for value in a set
  20     template <typename T, int N, typename U>
  21     inline bool contains(const T (&v)[N], U i)
  22     {
  23         return std::find_if(std::begin(v), std::end(v),
  24                 [=](T p){ return i == static_cast<U>(p); }) != std::end(v);
  25     }
  26 }
  27
  28 bool control_conn_t::process_packet()
  29 {
  30     using std::string;
  31
  32     // Note that where we call queue_packet, we must generally check the return value. If it
  33     // returns false it has either deleted the connection or marked it for deletion; we
  34     // shouldn't touch instance members after that point.
  35
  36     int pktType = rbuf[0];
  37     if (pktType == DINIT_CP_QUERYVERSION) {
  38         // Responds with:
  39         // DINIT_RP_CVERSION, (2 byte) minimum compatible version, (2 byte) actual version
  40         char replyBuf[] = { DINIT_RP_CPVERSION, 0, 0, 0, 0 };
  41         memcpy(replyBuf + 1, &min_compat_version, 2);
  42         memcpy(replyBuf + 3, &cp_version, 2);
  43         if (! queue_packet(replyBuf, sizeof(replyBuf))) return false;
  44         rbuf.consume(1);
  45         return true;
  46     }
  47     if (pktType == DINIT_CP_FINDSERVICE || pktType == DINIT_CP_LOADSERVICE) {
  48         return process_find_load(pktType);
  49     }
  50     if (pktType == DINIT_CP_STARTSERVICE || pktType == DINIT_CP_STOPSERVICE
  51             || pktType == DINIT_CP_WAKESERVICE || pktType == DINIT_CP_RELEASESERVICE) {
  52         return process_start_stop(pktType);
  53     }
  54     if (pktType == DINIT_CP_UNPINSERVICE) {
  55         return process_unpin_service();
  56     }
  57     if (pktType == DINIT_CP_UNLOADSERVICE) {
  58         return process_unload_service();
  59     }
  60     if (pktType == DINIT_CP_SHUTDOWN) {
  61         // Shutdown/reboot
  62         if (rbuf.get_length() < 2) {
  63             chklen = 2;
  64             return true;
  65         }
  66
  67         if (contains({shutdown_type_t::REMAIN, shutdown_type_t::HALT,
  68                 shutdown_type_t::POWEROFF, shutdown_type_t::REBOOT}, rbuf[1])) {
  69             auto sd_type = static_cast<shutdown_type_t>(rbuf[1]);
  70
  71             services->stop_all_services(sd_type);
  72             char ackBuf[] = { DINIT_RP_ACK };
  73             if (! queue_packet(ackBuf, 1)) return false;
  74
  75             // Clear the packet from the buffer
  76             rbuf.consume(2);
  77             chklen = 0;
  78             return true;
  79         }
  80
  81         // (otherwise fall through to below).
  82     }
  83     if (pktType == DINIT_CP_LISTSERVICES) {
  84         return list_services();
  85     }
  86     if (pktType == DINIT_CP_ADD_DEP) {
  87         return add_service_dep();
  88     }
  89     if (pktType == DINIT_CP_REM_DEP) {
  90         return rm_service_dep();
  91     }
  92     if (pktType == DINIT_CP_QUERY_LOAD_MECH) {
  93         return query_load_mech();
  94     }
  95     if (pktType == DINIT_CP_ENABLESERVICE) {
  96         return add_service_dep(true);
  97     }
  98     if (pktType == DINIT_CP_QUERYSERVICENAME) {
  99         return process_query_name();
 100     }
 101
 102     // Unrecognized: give error response
 103     char outbuf[] = { DINIT_RP_BADREQ };
 104     if (! queue_packet(outbuf, 1)) return false;
 105     bad_conn_close = true;
 106     iob.set_watches(OUT_EVENTS);
 107     return true;
 108 }
 109
 110 bool control_conn_t::process_find_load(int pktType)
 111 {
 112     using std::string;
 113
 114     constexpr int pkt_size = 4;
 115
 116     if (rbuf.get_length() < pkt_size) {
 117         chklen = pkt_size;
 118         return true;
 119     }
 120
 121     uint16_t svcSize;
 122     rbuf.extract((char *)&svcSize, 1, 2);
 123     if (svcSize <= 0 || svcSize > (1024 - 3)) {
 124         // Queue error response / mark connection bad
 125         char badreqRep[] = { DINIT_RP_BADREQ };
 126         if (! queue_packet(badreqRep, 1)) return false;
 127         bad_conn_close = true;
 128         iob.set_watches(OUT_EVENTS);
 129         return true;
 130     }
 131     chklen = svcSize + 3; // packet type + (2 byte) length + service name
 132
 133     if (rbuf.get_length() < chklen) {
 134         // packet not complete yet; read more
 135         return true;
 136     }
 137
 138     service_record * record = nullptr;
 139
 140     string serviceName = rbuf.extract_string(3, svcSize);
 141
 142     if (pktType == DINIT_CP_LOADSERVICE) {
 143         // LOADSERVICE
 144         try {
 145             record = services->load_service(serviceName.c_str());
 146         }
 147         catch (service_load_exc &slexc) {
 148             log(loglevel_t::ERROR, "Could not load service ", slexc.service_name, ": ",
 149                     slexc.exc_description);
 150         }
 151     }
 152     else {
 153         // FINDSERVICE
 154         record = services->find_service(serviceName.c_str());
 155     }
 156
 157     if (record != nullptr) {
 158         // Allocate a service handle
 159         handle_t handle = allocate_service_handle(record);
 160         std::vector<char> rp_buf;
 161         rp_buf.reserve(7);
 162         rp_buf.push_back(DINIT_RP_SERVICERECORD);
 163         rp_buf.push_back(static_cast<char>(record->get_state()));
 164         for (int i = 0; i < (int) sizeof(handle); i++) {
 165             rp_buf.push_back(*(((char *) &handle) + i));
 166         }
 167         rp_buf.push_back(static_cast<char>(record->get_target_state()));
 168         if (! queue_packet(std::move(rp_buf))) return false;
 169     }
 170     else {
 171         std::vector<char> rp_buf = { DINIT_RP_NOSERVICE };
 172         if (! queue_packet(std::move(rp_buf))) return false;
 173     }
 174
 175     // Clear the packet from the buffer
 176     rbuf.consume(chklen);
 177     chklen = 0;
 178     return true;
 179 }
 180
 181 bool control_conn_t::check_dependents(service_record *service, bool &had_dependents)
 182 {
 183     std::vector<char> reply_pkt;
 184     size_t num_depts = 0;
 185
 186     for (service_dep *dep : service->get_dependents()) {
 187         if (dep->dep_type == dependency_type::REGULAR) {
 188             num_depts++;
 189             // find or allocate a service handle
 190             handle_t dept_handle = allocate_service_handle(dep->get_from());
 191             if (reply_pkt.empty()) {
 192                 // packet type, size
 193                 reply_pkt.reserve(1 + sizeof(size_t) + sizeof(handle_t));
 194                 reply_pkt.resize(1 + sizeof(size_t));
 195                 reply_pkt[0] = DINIT_RP_DEPENDENTS;
 196             }
 197             auto old_size = reply_pkt.size();
 198             reply_pkt.resize(old_size + sizeof(handle_t));
 199             memcpy(reply_pkt.data() + old_size, &dept_handle, sizeof(dept_handle));
 200         }
 201     }
 202
 203     if (num_depts != 0) {
 204         // There are affected dependents
 205         had_dependents = true;
 206         memcpy(reply_pkt.data() + 1, &num_depts, sizeof(num_depts));
 207         return queue_packet(std::move(reply_pkt));
 208     }
 209
 210     had_dependents = false;
 211     return true;
 212 }
 213
 214 bool control_conn_t::process_start_stop(int pktType)
 215 {
 216     using std::string;
 217
 218     constexpr int pkt_size = 2 + sizeof(handle_t);
 219
 220     if (rbuf.get_length() < pkt_size) {
 221         chklen = pkt_size;
 222         return true;
 223     }
 224
 225     // 1 byte: packet type
 226     // 1 byte: pin in requested state (0 = no pin, 1 = pin)
 227     // 4 bytes: service handle
 228
 229     bool do_pin = ((rbuf[1] & 1) == 1);
 230     handle_t handle;
 231     rbuf.extract((char *) &handle, 2, sizeof(handle));
 232
 233     service_record *service = find_service_for_key(handle);
 234     if (service == nullptr) {
 235         // Service handle is bad
 236         char badreqRep[] = { DINIT_RP_BADREQ };
 237         if (! queue_packet(badreqRep, 1)) return false;
 238         bad_conn_close = true;
 239         iob.set_watches(OUT_EVENTS);
 240         return true;
 241     }
 242     else {
 243         char ack_buf[1] = { DINIT_RP_ACK };
 244
 245         switch (pktType) {
 246         case DINIT_CP_STARTSERVICE:
 247             // start service, mark as required
 248             if (services->is_shutting_down()) {
 249                 ack_buf[0] = DINIT_RP_NAK;
 250                 break;
 251             }
 252             if (do_pin) service->pin_start();
 253             service->start();
 254             services->process_queues();
 255             if (service->get_state() == service_state_t::STARTED) ack_buf[0] = DINIT_RP_ALREADYSS;
 256             break;
 257         case DINIT_CP_STOPSERVICE:
 258         {
 259             // force service to stop
 260             bool gentle = ((rbuf[1] & 2) == 2);
 261             if (gentle) {
 262                 // Check dependents; return appropriate response if any will be affected
 263                 bool has_dependents;
 264                 if (! check_dependents(service, has_dependents)) {
 265                     return false;
 266                 }
 267                 if (has_dependents) {
 268                     // Reply packet has already been sent
 269                     goto clear_out;
 270                 }
 271             }
 272             if (do_pin) service->pin_stop();
 273             service->stop(true);
 274             service->forced_stop();
 275             services->process_queues();
 276             if (service->get_state() == service_state_t::STOPPED) ack_buf[0] = DINIT_RP_ALREADYSS;
 277             break;
 278         }
 279         case DINIT_CP_WAKESERVICE:
 280             // re-start a stopped service (do not mark as required)
 281             if (services->is_shutting_down()) {
 282                 ack_buf[0] = DINIT_RP_NAK;
 283                 break;
 284             }
 285             if (do_pin) service->pin_start();
 286             service->start(false);
 287             services->process_queues();
 288             if (service->get_state() == service_state_t::STARTED) ack_buf[0] = DINIT_RP_ALREADYSS;
 289             break;
 290         case DINIT_CP_RELEASESERVICE:
 291             // remove required mark, stop if not required by dependents
 292             if (do_pin) service->pin_stop();
 293             service->stop(false);
 294             services->process_queues();
 295             if (service->get_state() == service_state_t::STOPPED) ack_buf[0] = DINIT_RP_ALREADYSS;
 296             break;
 297         }
 298
 299         if (! queue_packet(ack_buf, 1)) return false;
 300     }
 301
 302     clear_out:
 303     // Clear the packet from the buffer
 304     rbuf.consume(pkt_size);
 305     chklen = 0;
 306     return true;
 307 }
 308
 309 bool control_conn_t::process_unpin_service()
 310 {
 311     using std::string;
 312
 313     constexpr int pkt_size = 1 + sizeof(handle_t);
 314
 315     if (rbuf.get_length() < pkt_size) {
 316         chklen = pkt_size;
 317         return true;
 318     }
 319
 320     // 1 byte: packet type
 321     // 4 bytes: service handle
 322
 323     handle_t handle;
 324     rbuf.extract((char *) &handle, 1, sizeof(handle));
 325
 326     service_record *service = find_service_for_key(handle);
 327     if (service == nullptr) {
 328         // Service handle is bad
 329         char badreqRep[] = { DINIT_RP_BADREQ };
 330         if (! queue_packet(badreqRep, 1)) return false;
 331         bad_conn_close = true;
 332         iob.set_watches(OUT_EVENTS);
 333         return true;
 334     }
 335
 336     service->unpin();
 337     services->process_queues();
 338     char ack_buf[] = { (char) DINIT_RP_ACK };
 339     if (! queue_packet(ack_buf, 1)) return false;
 340
 341     // Clear the packet from the buffer
 342     rbuf.consume(pkt_size);
 343     chklen = 0;
 344     return true;
 345 }
 346
 347 bool control_conn_t::process_unload_service()
 348 {
 349     using std::string;
 350
 351     constexpr int pkt_size = 1 + sizeof(handle_t);
 352
 353     if (rbuf.get_length() < pkt_size) {
 354         chklen = pkt_size;
 355         return true;
 356     }
 357
 358     // 1 byte: packet type
 359     // 4 bytes: service handle
 360
 361     handle_t handle;
 362     rbuf.extract((char *) &handle, 1, sizeof(handle));
 363
 364     service_record *service = find_service_for_key(handle);
 365     if (service == nullptr) {
 366         // Service handle is bad
 367         char badreq_rep[] = { DINIT_RP_BADREQ };
 368         if (! queue_packet(badreq_rep, 1)) return false;
 369         bad_conn_close = true;
 370         iob.set_watches(OUT_EVENTS);
 371         return true;
 372     }
 373
 374     if (! service->has_lone_ref() || service->get_state() != service_state_t::STOPPED) {
 375         // Cannot unload: has other references
 376         char nak_rep[] = { DINIT_RP_NAK };
 377         if (! queue_packet(nak_rep, 1)) return false;
 378     }
 379     else {
 380         // unload
 381         service->prepare_for_unload();
 382         services->remove_service(service);
 383         delete service;
 384
 385         // drop handle
 386         service_key_map.erase(service);
 387         key_service_map.erase(handle);
 388
 389         // send ack
 390         char ack_buf[] = { (char) DINIT_RP_ACK };
 391         if (! queue_packet(ack_buf, 1)) return false;
 392     }
 393
 394     // Clear the packet from the buffer
 395     rbuf.consume(pkt_size);
 396     chklen = 0;
 397     return true;
 398 }
 399
 400 bool control_conn_t::list_services()
 401 {
 402     rbuf.consume(1); // clear request packet
 403     chklen = 0;
 404
 405     try {
 406         auto slist = services->list_services();
 407         for (auto sptr : slist) {
 408             std::vector<char> pkt_buf;
 409
 410             int hdrsize = 8 + std::max(sizeof(int), sizeof(pid_t));
 411
 412             const std::string &name = sptr->get_name();
 413             int nameLen = std::min((size_t)256, name.length());
 414             pkt_buf.resize(hdrsize + nameLen);
 415
 416             pkt_buf[0] = DINIT_RP_SVCINFO;
 417             pkt_buf[1] = nameLen;
 418             pkt_buf[2] = static_cast<char>(sptr->get_state());
 419             pkt_buf[3] = static_cast<char>(sptr->get_target_state());
 420
 421             char b0 = sptr->is_waiting_for_console() ? 1 : 0;
 422             b0 |= sptr->has_console() ? 2 : 0;
 423             b0 |= sptr->was_start_skipped() ? 4 : 0;
 424             pkt_buf[4] = b0;
 425             pkt_buf[5] = static_cast<char>(sptr->get_stop_reason());
 426
 427             pkt_buf[6] = 0; // reserved
 428             pkt_buf[7] = 0;
 429
 430             // Next: either the exit status, or the process ID
 431             if (sptr->get_state() != service_state_t::STOPPED) {
 432                 pid_t proc_pid = sptr->get_pid();
 433                 memcpy(pkt_buf.data() + 8, &proc_pid, sizeof(proc_pid));
 434             }
 435             else {
 436                 int exit_status = sptr->get_exit_status();
 437                 memcpy(pkt_buf.data() + 8, &exit_status, sizeof(exit_status));
 438             }
 439
 440             for (int i = 0; i < nameLen; i++) {
 441                 pkt_buf[hdrsize+i] = name[i];
 442             }
 443
 444             if (! queue_packet(std::move(pkt_buf))) return false;
 445         }
 446
 447         char ack_buf[] = { (char) DINIT_RP_LISTDONE };
 448         if (! queue_packet(ack_buf, 1)) return false;
 449
 450         return true;
 451     }
 452     catch (std::bad_alloc &exc)
 453     {
 454         do_oom_close();
 455         return true;
 456     }
 457 }
 458
 459 bool control_conn_t::add_service_dep(bool do_enable)
 460 {
 461     // 1 byte packet type
 462     // 1 byte dependency type
 463     // handle: "from"
 464     // handle: "to"
 465
 466     constexpr int pkt_size = 2 + sizeof(handle_t) * 2;
 467
 468     if (rbuf.get_length() < pkt_size) {
 469         chklen = pkt_size;
 470         return true;
 471     }
 472
 473     handle_t from_handle;
 474     handle_t to_handle;
 475     rbuf.extract((char *) &from_handle, 2, sizeof(from_handle));
 476     rbuf.extract((char *) &to_handle, 2 + sizeof(from_handle), sizeof(to_handle));
 477
 478     service_record *from_service = find_service_for_key(from_handle);
 479     service_record *to_service = find_service_for_key(to_handle);
 480     if (from_service == nullptr || to_service == nullptr || from_service == to_service) {
 481         // Service handle is bad
 482         char badreq_rep[] = { DINIT_RP_BADREQ };
 483         if (! queue_packet(badreq_rep, 1)) return false;
 484         bad_conn_close = true;
 485         iob.set_watches(OUT_EVENTS);
 486         return true;
 487     }
 488
 489     // Check dependency type is valid:
 490     int dep_type_int = rbuf[1];
 491     if (! contains({dependency_type::MILESTONE, dependency_type::REGULAR,
 492             dependency_type::WAITS_FOR}, dep_type_int)) {
 493         char badreqRep[] = { DINIT_RP_BADREQ };
 494         if (! queue_packet(badreqRep, 1)) return false;
 495         bad_conn_close = true;
 496         iob.set_watches(OUT_EVENTS);
 497     }
 498     dependency_type dep_type = static_cast<dependency_type>(dep_type_int);
 499
 500     // Check current service states are valid for given dep type
 501     if (dep_type == dependency_type::REGULAR) {
 502         if (from_service->get_state() != service_state_t::STOPPED &&
 503                 to_service->get_state() != service_state_t::STARTED) {
 504             // Cannot create dependency now since it would be contradicted:
 505             char nak_rep[] = { DINIT_RP_NAK };
 506             if (! queue_packet(nak_rep, 1)) return false;
 507             rbuf.consume(pkt_size);
 508             chklen = 0;
 509             return true;
 510         }
 511     }
 512
 513     // Check for creation of circular dependency chain
 514     std::unordered_set<service_record *> dep_marks;
 515     std::vector<service_record *> dep_queue;
 516     dep_queue.push_back(to_service);
 517     while (! dep_queue.empty()) {
 518         service_record * sr = dep_queue.back();
 519         dep_queue.pop_back();
 520         // iterate deps; if dep == from, abort; otherwise add to set/queue
 521         // (only add to queue if not already in set)
 522         for (auto &dep : sr->get_dependencies()) {
 523             service_record * dep_to = dep.get_to();
 524             if (dep_to == from_service) {
 525                 // fail, circular dependency!
 526                 char nak_rep[] = { DINIT_RP_NAK };
 527                 if (! queue_packet(nak_rep, 1)) return false;
 528                 rbuf.consume(pkt_size);
 529                 chklen = 0;
 530                 return true;
 531             }
 532             if (dep_marks.insert(dep_to).second) {
 533                 dep_queue.push_back(dep_to);
 534             }
 535         }
 536     }
 537     dep_marks.clear();
 538     dep_queue.clear();
 539
 540     bool dep_exists = false;
 541     service_dep * dep_record = nullptr;
 542
 543     // Prevent creation of duplicate dependency:
 544     for (auto &dep : from_service->get_dependencies()) {
 545         service_record * dep_to = dep.get_to();
 546         if (dep_to == to_service && dep.dep_type == dep_type) {
 547             // Dependency already exists
 548             dep_exists = true;
 549             dep_record = &dep;
 550             break;
 551         }
 552     }
 553
 554     if (! dep_exists) {
 555         // Create dependency:
 556         dep_record = &(from_service->add_dep(to_service, dep_type));
 557         services->process_queues();
 558     }
 559
 560     if (do_enable && contains({service_state_t::STARTED, service_state_t::STARTING},
 561             from_service->get_state())) {
 562         // The dependency record is activated: mark it as holding acquisition of the dependency, and start
 563         // the dependency.
 564         if (!services->is_shutting_down()) {
 565             dep_record->get_from()->start_dep(*dep_record);
 566             services->process_queues();
 567         }
 568     }
 569
 570     char ack_rep[] = { DINIT_RP_ACK };
 571     if (! queue_packet(ack_rep, 1)) return false;
 572     rbuf.consume(pkt_size);
 573     chklen = 0;
 574     return true;
 575 }
 576
 577 bool control_conn_t::rm_service_dep()
 578 {
 579     // 1 byte packet type
 580     // 1 byte dependency type
 581     // handle: "from"
 582     // handle: "to"
 583
 584     constexpr int pkt_size = 2 + sizeof(handle_t) * 2;
 585
 586     if (rbuf.get_length() < pkt_size) {
 587         chklen = pkt_size;
 588         return true;
 589     }
 590
 591     handle_t from_handle;
 592     handle_t to_handle;
 593     rbuf.extract((char *) &from_handle, 2, sizeof(from_handle));
 594     rbuf.extract((char *) &to_handle, 2 + sizeof(from_handle), sizeof(to_handle));
 595
 596     service_record *from_service = find_service_for_key(from_handle);
 597     service_record *to_service = find_service_for_key(to_handle);
 598     if (from_service == nullptr || to_service == nullptr || from_service == to_service) {
 599         // Service handle is bad
 600         char badreq_rep[] = { DINIT_RP_BADREQ };
 601         if (! queue_packet(badreq_rep, 1)) return false;
 602         bad_conn_close = true;
 603         iob.set_watches(OUT_EVENTS);
 604         return true;
 605     }
 606
 607     // Check dependency type is valid:
 608     int dep_type_int = rbuf[1];
 609     if (! contains({dependency_type::MILESTONE, dependency_type::REGULAR,
 610             dependency_type::WAITS_FOR}, dep_type_int)) {
 611         char badreqRep[] = { DINIT_RP_BADREQ };
 612         if (! queue_packet(badreqRep, 1)) return false;
 613         bad_conn_close = true;
 614         iob.set_watches(OUT_EVENTS);
 615     }
 616     dependency_type dep_type = static_cast<dependency_type>(dep_type_int);
 617
 618     // Remove dependency:
 619     from_service->rm_dep(to_service, dep_type);
 620     services->process_queues();
 621
 622     char ack_rep[] = { DINIT_RP_ACK };
 623     if (! queue_packet(ack_rep, 1)) return false;
 624     rbuf.consume(pkt_size);
 625     chklen = 0;
 626     return true;
 627 }
 628
 629 bool control_conn_t::process_query_name()
 630 {
 631     // 1 byte packet type
 632     // 1 byte reserved
 633     // handle: service
 634     constexpr int pkt_size = 2 + sizeof(handle_t);
 635
 636     if (rbuf.get_length() < pkt_size) {
 637         chklen = pkt_size;
 638         return true;
 639     }
 640
 641     // Reply:
 642     // 1 byte packet type = DINIT_RP_SERVICENAME
 643     // 1 byte reserved
 644     // uint16_t length
 645     // N bytes name
 646
 647     handle_t handle;
 648     rbuf.extract(&handle, 2, sizeof(handle));
 649     rbuf.consume(pkt_size);
 650     chklen = 0;
 651
 652     service_record *service = find_service_for_key(handle);
 653     if (service == nullptr || service->get_name().length() > std::numeric_limits<uint16_t>::max()) {
 654         char ack_rep[] = { DINIT_RP_ACK };
 655         if (! queue_packet(ack_rep, 1)) return false;
 656     }
 657
 658     std::vector<char> reply;
 659     const std::string &name = service->get_name();
 660     uint16_t name_length = name.length();
 661     reply.resize(2 + sizeof(uint16_t) + name_length);
 662     reply[0] = DINIT_RP_SERVICENAME;
 663     memcpy(reply.data() + 2, &name_length, sizeof(name_length));
 664     memcpy(reply.data() + 2 + sizeof(uint16_t), name.c_str(), name_length);
 665
 666     return queue_packet(std::move(reply));
 667 }
 668
 669 bool control_conn_t::query_load_mech()
 670 {
 671     rbuf.consume(1);
 672     chklen = 0;
 673
 674     if (services->get_set_type_id() == SSET_TYPE_DIRLOAD) {
 675         dirload_service_set *dss = static_cast<dirload_service_set *>(services);
 676         std::vector<char> reppkt;
 677         reppkt.resize(2 + sizeof(uint32_t) * 2);  // packet type, loader type, packet size, # dirs
 678         reppkt[0] = DINIT_RP_LOADER_MECH;
 679         reppkt[1] = SSET_TYPE_DIRLOAD;
 680
 681         // Number of directories in load path:
 682         uint32_t sdirs = dss->get_service_dir_count();
 683         std::memcpy(reppkt.data() + 2 + sizeof(uint32_t), &sdirs, sizeof(sdirs));
 684
 685         // Our current working directory, which above are relative to:
 686         // leave sizeof(uint32_t) for size, which we'll fill in afterwards:
 687         std::size_t curpos = reppkt.size() + sizeof(uint32_t);
 688 #ifdef PATH_MAX
 689         uint32_t try_path_size = PATH_MAX;
 690 #else
 691         uint32_t try_path_size = 2048;
 692 #endif
 693         char *wd;
 694         while (true) {
 695             std::size_t total_size = curpos + std::size_t(try_path_size);
 696             if (total_size < curpos) {
 697                 // Overflow. In theory we could now limit to size_t max, but the size must already
 698                 // be crazy long; let's abort.
 699                 char ack_rep[] = { DINIT_RP_NAK };
 700                 if (! queue_packet(ack_rep, 1)) return false;
 701                 return true;
 702             }
 703             reppkt.resize(total_size);
 704             wd = getcwd(reppkt.data() + curpos, try_path_size);
 705             if (wd != nullptr) break;
 706
 707             // Keep doubling the path size we try until it's big enough, or we get numeric overflow
 708             uint32_t new_try_path_size = try_path_size * uint32_t(2u);
 709             if (new_try_path_size < try_path_size) {
 710                 // Overflow.
 711                 char ack_rep[] = { DINIT_RP_NAK };
 712                 return queue_packet(ack_rep, 1);
 713             }
 714             try_path_size = new_try_path_size;
 715         }
 716
 717         uint32_t wd_len = std::strlen(reppkt.data() + curpos);
 718         reppkt.resize(curpos + std::size_t(wd_len));
 719         std::memcpy(reppkt.data() + curpos - sizeof(uint32_t), &wd_len, sizeof(wd_len));
 720
 721         // Each directory in the load path:
 722         for (int i = 0; uint32_t(i) < sdirs; i++) {
 723             const char *sdir = dss->get_service_dir(i);
 724             uint32_t dlen = std::strlen(sdir);
 725             auto cursize = reppkt.size();
 726             reppkt.resize(cursize + sizeof(dlen) + dlen);
 727             std::memcpy(reppkt.data() + cursize, &dlen, sizeof(dlen));
 728             std::memcpy(reppkt.data() + cursize + sizeof(dlen), sdir, dlen);
 729         }
 730
 731         // Total packet size:
 732         uint32_t fsize = reppkt.size();
 733         std::memcpy(reppkt.data() + 2, &fsize, sizeof(fsize));
 734
 735         if (! queue_packet(std::move(reppkt))) return false;
 736         return true;
 737     }
 738     else {
 739         // If we don't know how to deal with the service set type, send a NAK reply:
 740         char ack_rep[] = { DINIT_RP_NAK };
 741         return queue_packet(ack_rep, 1);
 742     }
 743 }
 744
 745 control_conn_t::handle_t control_conn_t::allocate_service_handle(service_record *record)
 746 {
 747     // Try to find a unique handle (integer) in a single pass. Since the map is ordered, we can search until
 748     // we find a gap in the handle values.
 749     handle_t candidate = 0;
 750     for (auto p : key_service_map) {
 751         if (p.first == candidate) ++candidate;
 752         else break;
 753     }
 754
 755     bool is_unique = (service_key_map.find(record) == service_key_map.end());
 756
 757     // The following operations perform allocation (can throw std::bad_alloc). If an exception occurs we
 758     // must undo any previous actions:
 759     if (is_unique) {
 760         record->add_listener(this);
 761     }
 762
 763     try {
 764         key_service_map[candidate] = record;
 765         service_key_map.insert(std::make_pair(record, candidate));
 766     }
 767     catch (...) {
 768         if (is_unique) {
 769             record->remove_listener(this);
 770         }
 771
 772         key_service_map.erase(candidate);
 773     }
 774
 775     return candidate;
 776 }
 777
 778 bool control_conn_t::queue_packet(const char *pkt, unsigned size) noexcept
 779 {
 780     int in_flag = bad_conn_close ? 0 : IN_EVENTS;
 781     bool was_empty = outbuf.empty();
 782
 783     // If the queue is empty, we can try to write the packet out now rather than queueing it.
 784     // If the write is unsuccessful or partial, we queue the remainder.
 785     if (was_empty) {
 786         int wr = bp_sys::write(iob.get_watched_fd(), pkt, size);
 787         if (wr == -1) {
 788             if (errno == EPIPE) {
 789                 return false;
 790             }
 791             if (errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) {
 792                 log(loglevel_t::WARN, "Error writing to control connection: ", strerror(errno));
 793                 return false;
 794             }
 795             // EAGAIN etc: fall through to below
 796         }
 797         else {
 798             if ((unsigned)wr == size) {
 799                 // Ok, all written.
 800                 iob.set_watches(in_flag);
 801                 return true;
 802             }
 803             pkt += wr;
 804             size -= wr;
 805         }
 806     }
 807
 808     // Create a vector out of the (remaining part of the) packet:
 809     try {
 810         outbuf.emplace_back(pkt, pkt + size);
 811         iob.set_watches(in_flag | OUT_EVENTS);
 812         return true;
 813     }
 814     catch (std::bad_alloc &baexc) {
 815         // Mark the connection bad, and stop reading further requests
 816         bad_conn_close = true;
 817         oom_close = true;
 818         if (was_empty) {
 819             // We can't send out-of-memory response as we already wrote as much as we
 820             // could above. Neither can we later send the response since we have currently
 821             // sent an incomplete packet. All we can do is close the connection.
 822             return false;
 823         }
 824         else {
 825             iob.set_watches(OUT_EVENTS);
 826             return true;
 827         }
 828     }
 829 }
 830
 831 // This queue_packet method is frustratingly similar to the one above, but the subtle differences
 832 // make them extraordinary difficult to combine into a single method.
 833 bool control_conn_t::queue_packet(std::vector<char> &&pkt) noexcept
 834 {
 835     int in_flag = bad_conn_close ? 0 : IN_EVENTS;
 836     bool was_empty = outbuf.empty();
 837
 838     if (was_empty) {
 839         outpkt_index = 0;
 840         // We can try sending the packet immediately:
 841         int wr = bp_sys::write(iob.get_watched_fd(), pkt.data(), pkt.size());
 842         if (wr == -1) {
 843             if (errno == EPIPE) {
 844                 return false;
 845             }
 846             if (errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) {
 847                 log(loglevel_t::WARN, "Error writing to control connection: ", strerror(errno));
 848                 return false;
 849             }
 850             // EAGAIN etc: fall through to below
 851         }
 852         else {
 853             if ((unsigned)wr == pkt.size()) {
 854                 // Ok, all written.
 855                 iob.set_watches(in_flag);
 856                 return true;
 857             }
 858             outpkt_index = wr;
 859         }
 860     }
 861
 862     try {
 863         outbuf.emplace_back(pkt);
 864         iob.set_watches(in_flag | OUT_EVENTS);
 865         return true;
 866     }
 867     catch (std::bad_alloc &baexc) {
 868         // Mark the connection bad, and stop reading further requests
 869         bad_conn_close = true;
 870         oom_close = true;
 871         if (was_empty) {
 872             // We can't send out-of-memory response as we already wrote as much as we
 873             // could above. Neither can we later send the response since we have currently
 874             // sent an incomplete packet. All we can do is close the connection.
 875             return false;
 876         }
 877         else {
 878             iob.set_watches(OUT_EVENTS);
 879             return true;
 880         }
 881     }
 882 }
 883
 884 bool control_conn_t::data_ready() noexcept
 885 {
 886     int fd = iob.get_watched_fd();
 887
 888     int r = rbuf.fill(fd);
 889
 890     // Note file descriptor is non-blocking
 891     if (r == -1) {
 892         if (errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) {
 893             log(loglevel_t::WARN, "Error writing to control connection: ", strerror(errno));
 894             return true;
 895         }
 896         return false;
 897     }
 898
 899     if (r == 0) {
 900         return true;
 901     }
 902
 903     // complete packet?
 904     if (rbuf.get_length() >= chklen) {
 905         try {
 906             return !process_packet();
 907         }
 908         catch (std::bad_alloc &baexc) {
 909             do_oom_close();
 910             return false;
 911         }
 912     }
 913     else if (rbuf.get_length() == 1024) {
 914         // Too big packet
 915         log(loglevel_t::WARN, "Received too-large control package; dropping connection");
 916         bad_conn_close = true;
 917         iob.set_watches(OUT_EVENTS);
 918     }
 919     else {
 920         int out_flags = (bad_conn_close || !outbuf.empty()) ? OUT_EVENTS : 0;
 921         iob.set_watches(IN_EVENTS | out_flags);
 922     }
 923
 924     return false;
 925 }
 926
 927 bool control_conn_t::send_data() noexcept
 928 {
 929     if (outbuf.empty() && bad_conn_close) {
 930         if (oom_close) {
 931             // Send oom response
 932             char oomBuf[] = { DINIT_RP_OOM };
 933             bp_sys::write(iob.get_watched_fd(), oomBuf, 1);
 934         }
 935         return true;
 936     }
 937
 938     vector<char> & pkt = outbuf.front();
 939     char *data = pkt.data();
 940     int written = bp_sys::write(iob.get_watched_fd(), data + outpkt_index, pkt.size() - outpkt_index);
 941     if (written == -1) {
 942         if (errno == EPIPE) {
 943             // read end closed
 944             return true;
 945         }
 946         else if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
 947             // spurious readiness notification?
 948         }
 949         else {
 950             log(loglevel_t::ERROR, "Error writing to control connection: ", strerror(errno));
 951             return true;
 952         }
 953         return false;
 954     }
 955
 956     outpkt_index += written;
 957     if (outpkt_index == pkt.size()) {
 958         // We've finished this packet, move on to the next:
 959         outbuf.pop_front();
 960         outpkt_index = 0;
 961         if (outbuf.empty() && ! oom_close) {
 962             if (! bad_conn_close) {
 963                 iob.set_watches(IN_EVENTS);
 964             }
 965             else {
 966                 return true;
 967             }
 968         }
 969     }
 970
 971     return false;
 972 }
 973
 974 control_conn_t::~control_conn_t() noexcept
 975 {
 976     bp_sys::close(iob.get_watched_fd());
 977     iob.deregister(loop);
 978
 979     // Clear service listeners
 980     for (auto p : service_key_map) {
 981         p.first->remove_listener(this);
 982     }
 983
 984     active_control_conns--;
 985 }