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