Fix missing #include (needed for compiling with GCC 10)
[oweals/dinit.git] / src / shutdown.cc
index c751fd086fe422a3e9beb877a7cc3abb37aa84c4..c82af94d7f2ea87f2d3215ae4ef8f572954921a1 100644 (file)
@@ -1,10 +1,10 @@
 #include <cstddef>
 #include <cstdio>
 #include <csignal>
-#include <unistd.h>
 #include <cstring>
 #include <string>
 #include <iostream>
+#include <exception>
 
 #include <sys/reboot.h>
 #include <sys/types.h>
 #include <sys/un.h>
 #include <sys/wait.h>
 #include <sys/stat.h>
+#include <unistd.h>
 #include <fcntl.h>
 
 #include "cpbuffer.h"
 #include "control-cmds.h"
 #include "service-constants.h"
+#include "dinit-client.h"
+#include "dinit-util.h"
+#include "mconfig.h"
+
+#include "dasynq.h"
 
 // shutdown:  shut down the system
-// This utility communicates with the dinit daemon via a unix socket (/dev/initctl).
+// This utility communicates with the dinit daemon via a unix socket (specified in SYSCONTROLSOCKET).
+
+static constexpr uint16_t min_cp_version = 1;
+static constexpr uint16_t max_cp_version = 1;
+
+using loop_t = dasynq::event_loop_n;
+using rearm = dasynq::rearm;
+using clock_type = dasynq::clock_type;
+class subproc_buffer;
 
 void do_system_shutdown(shutdown_type_t shutdown_type);
-static void unmount_disks();
-static void swap_off();
-static void wait_for_reply(cpbuffer<1024> &rbuffer, int fd);
+static void unmount_disks(loop_t &loop, subproc_buffer &sub_buf);
+static void swap_off(loop_t &loop, subproc_buffer &sub_buf);
 
+constexpr static int subproc_bufsize = 4096;
 
-class ReadCPException
+constexpr static char output_lost_msg[] = "[Some output has not been shown due to buffer overflow]\n";
+
+// A buffer which maintains a series of overflow markers, used for capturing and echoing
+// subprocess output.
+class subproc_buffer : private cpbuffer<subproc_bufsize>
 {
+    using base = cpbuffer<subproc_bufsize>;
+
+    int overflow_marker = -1;
+    int last_overflow = -1;  // last marker in the series
+    const char *overflow_msg_ptr = nullptr;   // current position in overflow message
+    dasynq::event_loop_n &loop;
+    dasynq::event_loop_n::fd_watcher * out_watch;
+
     public:
-    int errcode;
-    ReadCPException(int err) : errcode(err) { }
+    enum class fill_status
+    {
+        OK,
+        ENDFILE,
+        FULL
+    };
+
+    subproc_buffer(dasynq::event_loop_n &loop_p, int out_fd) : loop(loop_p)
+    {
+        using loop_t = dasynq::event_loop_n;
+        using rearm = dasynq::rearm;
+
+        out_watch = loop_t::fd_watcher::add_watch(loop, out_fd, dasynq::OUT_EVENTS,
+                [&](loop_t &eloop, int fd, int flags) -> rearm {
+
+            auto fstatus = flush(STDOUT_FILENO);
+            if (fstatus == subproc_buffer::fill_status::ENDFILE) {
+                return rearm::DISARM;
+            }
+
+            return rearm::REARM;
+        });
+    }
+
+    ~subproc_buffer()
+    {
+        out_watch->deregister(loop);
+    }
+
+    // Fill buffer by reading from file descriptor. Note caller must set overflow marker
+    // if the buffer becomes full and more data is available.
+    fill_status fill(int fd)
+    {
+        int rem = get_free();
+
+        if (rem == 0) {
+            return fill_status::FULL;
+        }
+
+        int read = base::fill(fd, rem);
+        if (read <= 0) {
+            if (read == -1 && errno == EAGAIN) {
+                return fill_status::OK;
+            }
+            return fill_status::ENDFILE;
+        }
+
+        out_watch->set_enabled(loop, true);
+        return fill_status::OK;
+    }
+
+    // Append a message. If the message will not fit in the buffer, discard it and mark overflow.
+    void append(const char *msg)
+    {
+        out_watch->set_enabled(loop, true);
+        int len = strlen(msg);
+        if (subproc_bufsize - get_length() >= len) {
+            base::append(msg, len);
+        }
+        else {
+            mark_overflow();
+        }
+    }
+
+    // Append the given buffer, which must fit in the remaining space in this buffer.
+    void append(const char *buf, int len)
+    {
+        out_watch->set_enabled(loop, true);
+        base::append(buf, len);
+    }
+
+    int get_free()
+    {
+        return base::get_free();
+    }
+
+    // Write buffer contents out to file descriptor. The descriptor is assumed to be non-blocking.
+    // returns ENDFILE if there is no more content to flush (buffer is now empty) or OK otherwise.
+    fill_status flush(int fd)
+    {
+        int to_write = get_contiguous_length(get_ptr(0));
+
+        if (overflow_marker != -1) {
+            if (overflow_marker == 0) {
+                // output (remainder of) overflow message
+                int l = std::strlen(overflow_msg_ptr);
+                int r = write(fd, overflow_msg_ptr, l);
+                if (r == l) {
+                    // entire message has been written; next marker is in buffer
+                    int16_t overflow_marker16;
+                    extract(reinterpret_cast<char *>(&overflow_marker16), 0, sizeof(overflow_marker16));
+                    overflow_marker = overflow_marker16;
+                    consume(sizeof(overflow_marker16));
+
+                    // no more overflow markers?
+                    if (overflow_marker == -1) {
+                        last_overflow = -1;
+                    }
+                    return get_length() == 0 ? fill_status::ENDFILE : fill_status::OK;
+                }
+                if (r > 0) {
+                    overflow_msg_ptr += r;
+                }
+                return fill_status::OK;
+            }
+
+            to_write = std::min(to_write, overflow_marker);
+        }
+
+        int r = write(fd, get_ptr(0), to_write);
+        if (r > 0) {
+            consume(r);
+            if (overflow_marker != -1) {
+                overflow_marker -= r;
+                last_overflow -= r;
+                if (overflow_marker == 0) {
+                    overflow_msg_ptr = output_lost_msg;
+                }
+            }
+        }
+
+        return get_length() == 0 ? fill_status::ENDFILE : fill_status::OK;
+    }
+
+    // Mark overflow occurred. Call this only when the buffer is full.
+    // The marker is put after the most recent newline in the buffer, if possible, so that whole
+    // lines are retained in the buffer. In some cases marking overflow will not add a new overflow
+    // marker but simply trim the buffer to an existing marker.
+    void mark_overflow()
+    {
+        // Try to find the last newline in the buffer
+        int begin = 0;
+        if (last_overflow != -1) {
+            begin = last_overflow + sizeof(int16_t);
+        }
+        int end = get_length() - 1 - sizeof(int16_t); // -1, then -2 for storage of marker
+
+        int i;
+        for (i = end; i >= begin; i--) {
+            if ((*this)[i] == '\n') break;
+        }
+
+        if (last_overflow != -1 && i < begin) {
+            // No new line after existing marker: trim all beyond that marker, don't
+            // create a new marker:
+            trim_to(last_overflow + sizeof(uint16_t));
+            return;
+        }
+
+        if (i < begin) {
+            // No newline in the whole buffer... we'll put the overflow marker at the end,
+            // on the assumption that it is better to output a partial line than it is to
+            // discard the entire buffer:
+            last_overflow = get_length() - sizeof(int16_t);
+            overflow_marker = last_overflow;
+            int16_t overflow16 = -1;
+            char * overflow16_ptr = reinterpret_cast<char *>(&overflow16);
+            *get_ptr(last_overflow + 0) = overflow16_ptr[0];
+            *get_ptr(last_overflow + 1) = overflow16_ptr[1];
+            return;
+        }
+
+        // We found a newline, put the overflow marker just after it:
+        int new_overflow = i + 1;
+        if (last_overflow != -1) {
+            int16_t new_overflow16 = new_overflow;
+            char * new_overflow16_ptr = reinterpret_cast<char *>(&new_overflow16);
+            *get_ptr(last_overflow + 0) = new_overflow16_ptr[0];
+            *get_ptr(last_overflow + 1) = new_overflow16_ptr[1];
+        }
+        last_overflow = new_overflow;
+        if (overflow_marker == -1) {
+            overflow_marker = last_overflow;
+        }
+
+        int16_t overflow16 = -1;
+        char * overflow16_ptr = reinterpret_cast<char *>(&overflow16);
+        *get_ptr(last_overflow + 0) = overflow16_ptr[0];
+        *get_ptr(last_overflow + 0) = overflow16_ptr[1];
+        trim_to(last_overflow + sizeof(int16_t));
+    }
 };
 
 
@@ -44,6 +249,11 @@ int main(int argc, char **argv)
     bool use_passed_cfd = false;
     
     auto shutdown_type = shutdown_type_t::POWEROFF;
+
+    const char *execname = base_name(argv[0]);
+    if (strcmp(execname, "reboot") == 0) {
+        shutdown_type = shutdown_type_t::REBOOT;
+    }
         
     for (int i = 1; i < argc; i++) {
         if (argv[i][0] == '-') {
@@ -79,16 +289,16 @@ int main(int argc, char **argv)
     }
 
     if (show_help) {
-        cout << "dinit-shutdown :   shutdown the system" << endl;
-        cout << "  --help           : show this help" << endl;
-        cout << "  -r               : reboot" << endl;
-        cout << "  -h               : halt system" << endl;
-        cout << "  -p               : power down (default)" << endl;
-        cout << "  --use-passed-cfd : use the socket file descriptor identified by the DINIT_CS_FD" << endl;
-        cout << "                     environment variable to communicate with the init daemon." << endl;
-        cout << "  --system         : perform shutdown immediately, instead of issuing shutdown" << endl;
-        cout << "                     command to the init program. Not recommended for use" << endl;
-        cout << "                     by users." << endl;
+        cout << execname << " :   shutdown the system\n"
+                "  --help           : show this help\n"
+                "  -r               : reboot\n"
+                "  -h               : halt system\n"
+                "  -p               : power down (default)\n"
+                "  --use-passed-cfd : use the socket file descriptor identified by the DINIT_CS_FD\n"
+                "                     environment variable to communicate with the init daemon.\n"
+                "  --system         : perform shutdown immediately, instead of issuing shutdown\n"
+                "                     command to the init program. Not recommended for use\n"
+                "                     by users.\n";
         return 1;
     }
     
@@ -108,6 +318,9 @@ int main(int argc, char **argv)
             long int cfdnum = strtol(dinit_cs_fd_env, &endptr, 10);
             if (endptr != dinit_cs_fd_env) {
                 socknum = (int) cfdnum;
+                // Set non-blocking mode:
+                int sock_flags = fcntl(socknum, F_GETFL, 0);
+                fcntl(socknum, F_SETFL, sock_flags & ~O_NONBLOCK);
             }
             else {
                 use_passed_cfd = false;
@@ -125,7 +338,7 @@ int main(int argc, char **argv)
             return 1;
         }
         
-        const char *naddr = "/dev/dinitctl";
+        const char *naddr = SYSCONTROLSOCKET;
         
         struct sockaddr_un name;
         name.sun_family = AF_UNIX;
@@ -139,30 +352,24 @@ int main(int argc, char **argv)
         }
     }
 
-    // Build buffer;
-    constexpr int bufsize = 2;
-    char buf[bufsize];
-    
-    buf[0] = DINIT_CP_SHUTDOWN;
-    buf[1] = static_cast<char>(shutdown_type);
-    
-    cout << "Issuing shutdown command..." << endl;
-    
-    // TODO make sure to write the whole buffer
-    int r = write(socknum, buf, bufsize);
-    if (r == -1) {
-        perror("write");
-        return 1;
-    }
+    try {
+        cpbuffer_t rbuffer;
     
-    // Wait for ACK/NACK
-    // r = read(socknum, buf, 1);
-    //if (r > 0) {
-    //    cout << "Received acknowledgement. System should now shut down." << endl;
-    //}
+        check_protocol_version(min_cp_version, max_cp_version, rbuffer, socknum);
+
+        // Build buffer;
+        constexpr int bufsize = 2;
+        char buf[bufsize];
+
+        buf[0] = DINIT_CP_SHUTDOWN;
+        buf[1] = static_cast<char>(shutdown_type);
+
+        cout << "Issuing shutdown command..." << endl;
+
+        write_all_x(socknum, buf, bufsize);
+
+        // Wait for ACK/NACK
     
-    cpbuffer<1024> rbuffer;
-    try {
         wait_for_reply(rbuffer, socknum);
         
         if (rbuffer[0] != DINIT_RP_ACK) {
@@ -170,9 +377,20 @@ int main(int argc, char **argv)
             return 1;
         }
     }
-    catch (ReadCPException &exc)
-    {
-        cerr << "shutdown: control socket read failure or protocol error" << endl;    
+    catch (cp_old_client_exception &e) {
+        std::cerr << "shutdown: too old (server reports newer protocol version)" << std::endl;
+        return 1;
+    }
+    catch (cp_old_server_exception &e) {
+        std::cerr << "shutdown: server too old or protocol error" << std::endl;
+        return 1;
+    }
+    catch (cp_read_exception &e) {
+        cerr << "shutdown: control socket read failure or protocol error" << endl;
+        return 1;
+    }
+    catch (cp_write_exception &e) {
+        cerr << "shutdown: control socket write error: " << std::strerror(e.errcode) << endl;
         return 1;
     }
     
@@ -183,47 +401,6 @@ int main(int argc, char **argv)
     return 0;
 }
 
-// Fill a circular buffer from a file descriptor, reading at least _rlength_ bytes.
-// Throws ReadException if the requested number of bytes cannot be read, with:
-//     errcode = 0   if end of stream (remote end closed)
-//     errcode = errno   if another error occurred
-// Note that EINTR is ignored (i.e. the read will be re-tried).
-static void fillBufferTo(cpbuffer<1024> *buf, int fd, int rlength)
-{
-    do {
-        int r = buf->fill_to(fd, rlength);
-        if (r == -1) {
-            if (errno != EINTR) {
-                throw ReadCPException(errno);
-            }
-        }
-        else if (r == 0) {
-            throw ReadCPException(0);
-        }
-        else {
-            return;
-        }
-    }
-    while (true);
-}
-
-// Wait for a reply packet, skipping over any information packets
-// that are received in the meantime.
-static void wait_for_reply(cpbuffer<1024> &rbuffer, int fd)
-{
-    fillBufferTo(&rbuffer, fd, 1);
-    
-    while (rbuffer[0] >= 100) {
-        // Information packet; discard.
-        fillBufferTo(&rbuffer, fd, 2);
-        int pktlen = (unsigned char) rbuffer[1];
-        
-        rbuffer.consume(1);  // Consume one byte so we'll read one byte of the next packet
-        fillBufferTo(&rbuffer, fd, pktlen);
-        rbuffer.consume(pktlen - 1);
-    }
-}
-
 // Actually shut down the system.
 void do_system_shutdown(shutdown_type_t shutdown_type)
 {
@@ -234,10 +411,15 @@ void do_system_shutdown(shutdown_type_t shutdown_type)
     sigfillset(&allsigs);
     sigprocmask(SIG_SETMASK, &allsigs, nullptr);
     
-    int reboot_type = 0;
-    if (shutdown_type == shutdown_type_t::REBOOT) reboot_type = RB_AUTOBOOT;
-    else if (shutdown_type == shutdown_type_t::POWEROFF) reboot_type = RB_POWER_OFF;
-    else reboot_type = RB_HALT_SYSTEM;
+    int reboot_type = RB_AUTOBOOT; // reboot
+#if defined(RB_POWER_OFF)
+    if (shutdown_type == shutdown_type_t::POWEROFF) reboot_type = RB_POWER_OFF;
+#endif
+#if defined(RB_HALT_SYSTEM)
+    if (shutdown_type == shutdown_type_t::HALT) reboot_type = RB_HALT_SYSTEM;
+#elif defined(RB_HALT)
+    if (shutdown_type == shutdown_type_t::HALT) reboot_type = RB_HALT;
+#endif
     
     // Write to console rather than any terminal, since we lose the terminal it seems:
     close(STDOUT_FILENO);
@@ -246,48 +428,192 @@ void do_system_shutdown(shutdown_type_t shutdown_type)
         dup2(consfd, STDOUT_FILENO);
     }
     
-    cout << "Sending TERM/KILL to all processes..." << endl;
+    loop_t loop;
+    subproc_buffer sub_buf {loop, STDOUT_FILENO};
+
+    sub_buf.append("Sending TERM/KILL to all processes...\n");
     
     // Send TERM/KILL to all (remaining) processes
     kill(-1, SIGTERM);
-    sleep(1);
+
+    // 1 second delay (while outputting from sub_buf):
+    bool timeout_reached = false;
+    dasynq::time_val timeout {1, 0};
+    dasynq::time_val interval {0,0};
+    loop_t::timer::add_timer(loop, clock_type::MONOTONIC, true /* relative */,
+            timeout.get_timespec(), interval.get_timespec(),
+            [&](loop_t &eloop, int expiry_count) -> rearm {
+
+        timeout_reached = true;
+        return rearm::REMOVE;
+    });
+
+    do {
+      loop.run();
+    } while (! timeout_reached);
+
     kill(-1, SIGKILL);
     
     // perform shutdown
-    cout << "Turning off swap..." << endl;
-    swap_off();
-    cout << "Unmounting disks..." << endl;
-    unmount_disks();
+    sub_buf.append("Turning off swap...\n");
+    swap_off(loop, sub_buf);
+    sub_buf.append("Unmounting disks...\n");
+    unmount_disks(loop, sub_buf);
     sync();
     
-    cout << "Issuing shutdown via kernel..." << endl;
+    sub_buf.append("Issuing shutdown via kernel...\n");
+    loop.poll();  // give message a chance to get to console
     reboot(reboot_type);
 }
 
-static void unmount_disks()
+// Watcher for subprocess output.
+class subproc_out_watch : public loop_t::fd_watcher_impl<subproc_out_watch>
 {
-    pid_t chpid = fork();
-    if (chpid == 0) {
-        // umount -a -r
-        //  -a : all filesystems (except proc)
-        //  -r : mount readonly if can't unmount
-        execl("/bin/umount", "/bin/umount", "-a", "-r", nullptr);
+    subproc_buffer &sub_buf;
+    bool in_overflow = false;
+
+    rearm read_overflow(int fd)
+    {
+        char buf[128];
+        int r = read(fd, buf, 128);
+        if (r == 0 || (r == -1 && errno != EAGAIN)) {
+            return rearm::NOOP; // leave disarmed
+        }
+        if (r == -1) {
+            return rearm::REARM;
+        }
+
+        // How much space is available?
+        int fr = sub_buf.get_free();
+        for (int b = r - std::min(r, fr); b < r; b++) {
+            if (buf[b] == '\n') {
+                // Copy the (partial) line into sub_buf and leave overflow mode
+                sub_buf.append(buf + b, r - b);
+                in_overflow = false;
+            }
+        }
+        return rearm::REARM;
+    }
+
+    public:
+    subproc_out_watch(subproc_buffer &sub_buf_p) : sub_buf(sub_buf_p) {}
+
+    rearm fd_event(loop_t &, int fd, int flags)
+    {
+        // if current status is reading overflow, read and discard until newline
+        if (in_overflow) {
+            return read_overflow(fd);
+        }
+
+        auto r = sub_buf.fill(fd);
+        if (r == subproc_buffer::fill_status::FULL) {
+            sub_buf.mark_overflow();
+            in_overflow = true;
+            return read_overflow(fd);
+        }
+        else if (r == subproc_buffer::fill_status::ENDFILE) {
+            return rearm::NOOP;
+        }
+
+        return rearm::REARM;  // re-enable watcher
     }
-    else if (chpid > 0) {
-        int status;
-        waitpid(chpid, &status, 0);
+};
+
+// Run process, put its output through the subprocess buffer
+//   may throw: std::system_error, std::bad_alloc
+static void run_process(const char * prog_args[], loop_t &loop, subproc_buffer &sub_buf)
+{
+    class sp_watcher_t : public loop_t::child_proc_watcher_impl<sp_watcher_t>
+    {
+        public:
+        bool terminated = false;
+
+        rearm status_change(loop_t &, pid_t child, int status)
+        {
+            terminated = true;
+            return rearm::REMOVE;
+        }
+    };
+
+    sp_watcher_t sp_watcher;
+
+    // Create output pipe
+    bool have_pipe = true;
+    int pipefds[2];
+    if (dasynq::pipe2(pipefds, O_NONBLOCK) == -1) {
+        sub_buf.append("Warning: ");
+        sub_buf.append(prog_args[0]);
+        sub_buf.append(": could not create pipe for subprocess output\n");
+        have_pipe = false;
+        // Note, we proceed and let the sub-process run with our stdout/stderr.
+    }
+
+    subproc_out_watch owatch {sub_buf};
+
+    if (have_pipe) {
+        close(pipefds[1]);
+        try {
+            owatch.add_watch(loop, pipefds[0], dasynq::IN_EVENTS);
+        }
+        catch (...) {
+            // failed to create the watcher for the subprocess output; again, let it run with
+            // our stdout/stderr
+            sub_buf.append("Warning: could not create output watch for subprocess\n");
+            close(pipefds[0]);
+            have_pipe = false;
+        }
+    }
+
+    // If we've buffered any messages/output, give them a chance to go out now:
+    loop.poll();
+
+    pid_t ch_pid = sp_watcher.fork(loop);
+    if (ch_pid == 0) {
+        // child
+        // Dup output pipe to stdout, stderr
+        if (have_pipe) {
+            dup2(pipefds[1], STDOUT_FILENO);
+            dup2(pipefds[1], STDERR_FILENO);
+            close(pipefds[0]);
+            close(pipefds[1]);
+        }
+        execv(prog_args[0], const_cast<char **>(prog_args));
+        puts("Failed to execute subprocess: ");
+        perror(prog_args[0]);
+        _exit(1);
+    }
+
+    do {
+        loop.run();
+    } while (! sp_watcher.terminated);
+
+    if (have_pipe) {
+        owatch.deregister(loop);
     }
 }
 
-static void swap_off()
+static void unmount_disks(loop_t &loop, subproc_buffer &sub_buf)
 {
-    pid_t chpid = fork();
-    if (chpid == 0) {
-        // swapoff -a
-        execl("/sbin/swapoff", "/sbin/swapoff", "-a", nullptr);
+    try {
+        const char * unmount_args[] = { "/bin/umount", "-a", "-r", nullptr };
+        run_process(unmount_args, loop, sub_buf);
+    }
+    catch (std::exception &e) {
+        sub_buf.append("Couldn't fork for umount: ");
+        sub_buf.append(e.what());
+        sub_buf.append("\n");
+    }
+}
+
+static void swap_off(loop_t &loop, subproc_buffer &sub_buf)
+{
+    try {
+        const char * swapoff_args[] = { "/sbin/swapoff", "-a", nullptr };
+        run_process(swapoff_args, loop, sub_buf);
     }
-    else if (chpid > 0) {
-        int status;
-        waitpid(chpid, &status, 0);
+    catch (std::exception &e) {
+        sub_buf.append("Couldn't fork for swapoff: ");
+        sub_buf.append(e.what());
+        sub_buf.append("\n");
     }
 }