control.cc 29 KB

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