Handle fork+exec failure in a separate virtual method.
[oweals/dinit.git] / src / control.h
1 #ifndef DINIT_CONTROL_H
2 #define DINIT_CONTROL_H
3
4 #include <list>
5 #include <vector>
6 #include <unordered_map>
7 #include <limits>
8 #include <cstddef>
9
10 #include <unistd.h>
11
12 #include "dasynq.h"
13
14 #include "dinit-log.h"
15 #include "control-cmds.h"
16 #include "service-listener.h"
17 #include "cpbuffer.h"
18
19 // Control connection for dinit
20
21 using namespace dasynq;
22 using eventloop_t = event_loop<null_mutex>;
23
24 class control_conn_t;
25 class control_conn_watcher;
26
27 // forward-declaration of callback:
28 static rearm control_conn_cb(eventloop_t *loop, control_conn_watcher *watcher, int revents);
29
30 // Pointer to the control connection that is listening for rollback completion
31 extern control_conn_t * rollback_handler_conn;
32
33 extern int active_control_conns;
34
35 // "packet" format:
36 // (1 byte) packet type
37 // (N bytes) additional data (service name, etc)
38 //   for LOADSERVICE/FINDSERVICE:
39 //      (2 bytes) service name length
40 //      (M bytes) service name (without nul terminator)
41
42 // Information packet:
43 // (1 byte) packet type, >= 100
44 // (1 byte) packet length (including all fields)
45 //       N bytes: packet data (N = (length - 2))
46
47 class service_set;
48 class service_record;
49
50 class control_conn_watcher : public eventloop_t::bidi_fd_watcher_impl<control_conn_watcher>
51 {
52     inline rearm receive_event(eventloop_t &loop, int fd, int flags) noexcept;
53
54     eventloop_t * event_loop;
55
56     public:
57     control_conn_watcher(eventloop_t & event_loop_p) : event_loop(&event_loop_p)
58     {
59         // constructor
60     }
61
62     rearm read_ready(eventloop_t &loop, int fd) noexcept
63     {
64         return receive_event(loop, fd, IN_EVENTS);
65     }
66     
67     rearm write_ready(eventloop_t &loop, int fd) noexcept
68     {
69         return receive_event(loop, fd, OUT_EVENTS);
70     }
71
72     void set_watches(int flags)
73     {
74         eventloop_t::bidi_fd_watcher::set_watches(*event_loop, flags);
75     }
76 };
77
78 inline rearm control_conn_watcher::receive_event(eventloop_t &loop, int fd, int flags) noexcept
79 {
80     return control_conn_cb(&loop, this, flags);
81 }
82
83
84 class control_conn_t : private service_listener
85 {
86     friend rearm control_conn_cb(eventloop_t *loop, control_conn_watcher *watcher, int revents);
87     
88     control_conn_watcher iob;
89     eventloop_t &loop;
90     service_set *services;
91     
92     bool bad_conn_close = false; // close when finished output?
93     bool oom_close = false;      // send final 'out of memory' indicator
94
95     // The packet length before we need to re-check if the packet is complete.
96     // process_packet() will not be called until the packet reaches this size.
97     int chklen;
98     
99     // Receive buffer
100     cpbuffer<1024> rbuf;
101     
102     template <typename T> using list = std::list<T>;
103     template <typename T> using vector = std::vector<T>;
104     
105     // A mapping between service records and their associated numerical identifier used
106     // in communction
107     using handle_t = uint32_t;
108     std::unordered_multimap<service_record *, handle_t> serviceKeyMap;
109     std::unordered_map<handle_t, service_record *> keyServiceMap;
110     
111     // Buffer for outgoing packets. Each outgoing back is represented as a vector<char>.
112     list<vector<char>> outbuf;
113     // Current index within the first outgoing packet (all previous bytes have been sent).
114     unsigned outpkt_index = 0;
115     
116     // Queue a packet to be sent
117     //  Returns:  false if the packet could not be queued and a suitable error packet
118     //              could not be sent/queued (the connection should be closed);
119     //            true (with bad_conn_close == false) if the packet was successfully
120     //              queued;
121     //            true (with bad_conn_close == true) if the packet was not successfully
122     //              queued (but a suitable error packate has been queued).
123     // The in/out watch enabled state will also be set appropriately.
124     bool queue_packet(vector<char> &&v) noexcept;
125     bool queue_packet(const char *pkt, unsigned size) noexcept;
126
127     // Process a packet.
128     //  Returns:  true (with bad_conn_close == false) if successful
129     //            true (with bad_conn_close == true) if an error packet was queued
130     //            false if an error occurred but no error packet could be queued
131     //                (connection should be closed).
132     // Throws:
133     //    std::bad_alloc - if an out-of-memory condition prevents processing
134     bool process_packet();
135     
136     // Process a STARTSERVICE/STOPSERVICE packet. May throw std::bad_alloc.
137     bool process_start_stop(int pktType);
138     
139     // Process a FINDSERVICE/LOADSERVICE packet. May throw std::bad_alloc.
140     bool process_find_load(int pktType);
141
142     // Process an UNPINSERVICE packet. May throw std::bad_alloc.
143     bool process_unpin_service();
144     
145     bool list_services();
146
147     // Notify that data is ready to be read from the socket. Returns true if the connection should
148     // be closed.
149     bool data_ready() noexcept;
150     
151     bool send_data() noexcept;
152     
153     // Allocate a new handle for a service; may throw std::bad_alloc
154     handle_t allocate_service_handle(service_record *record);
155     
156     service_record *find_service_for_key(uint32_t key)
157     {
158         try {
159             return keyServiceMap.at(key);
160         }
161         catch (std::out_of_range &exc) {
162             return nullptr;
163         }
164     }
165     
166     // Close connection due to out-of-memory condition.
167     void do_oom_close()
168     {
169         bad_conn_close = true;
170         oom_close = true;
171         iob.set_watches(OUT_EVENTS);
172     }
173     
174     // Process service event broadcast.
175     // Note that this can potentially be called during packet processing (upon issuing
176     // service start/stop orders etc).
177     void service_event(service_record * service, service_event_t event) noexcept final override
178     {
179         // For each service handle corresponding to the event, send an information packet.
180         auto range = serviceKeyMap.equal_range(service);
181         auto & i = range.first;
182         auto & end = range.second;
183         try {
184             while (i != end) {
185                 uint32_t key = i->second;
186                 std::vector<char> pkt;
187                 constexpr int pktsize = 3 + sizeof(key);
188                 pkt.reserve(pktsize);
189                 pkt.push_back(DINIT_IP_SERVICEEVENT);
190                 pkt.push_back(pktsize);
191                 char * p = (char *) &key;
192                 for (int j = 0; j < (int)sizeof(key); j++) {
193                     pkt.push_back(*p++);
194                 }
195                 pkt.push_back(static_cast<char>(event));
196                 queue_packet(std::move(pkt));
197                 ++i;
198             }
199         }
200         catch (std::bad_alloc &exc) {
201             do_oom_close();
202         }
203     }
204     
205     public:
206     control_conn_t(eventloop_t &loop, service_set * services_p, int fd)
207             : iob(loop), loop(loop), services(services_p), chklen(0)
208     {
209         iob.add_watch(loop, fd, IN_EVENTS);
210         active_control_conns++;
211     }
212     
213     bool rollback_complete() noexcept;
214         
215     virtual ~control_conn_t() noexcept;
216 };
217
218
219 static rearm control_conn_cb(eventloop_t * loop, control_conn_watcher * watcher, int revents)
220 {
221     char * cc_addr = (reinterpret_cast<char *>(watcher)) - offsetof(control_conn_t, iob);
222     control_conn_t *conn = reinterpret_cast<control_conn_t *>(cc_addr);
223     if (revents & IN_EVENTS) {
224         if (conn->data_ready()) {
225             delete conn;
226             return rearm::REMOVED;
227         }
228     }
229     if (revents & OUT_EVENTS) {
230         if (conn->send_data()) {
231             delete conn;
232             return rearm::REMOVED;
233         }
234     }
235     
236     return rearm::NOOP;
237 }
238
239 #endif