Properly exit when wayland server diesHEADmaster

2 files changed, 129 insertions, 37 deletions

diff --git a/Makefile b/Makefile
index d1eb8b1..5010ba0 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-CFLAGS=-g -Wall `pkg-config --cflags --libs libudev` 
+CFLAGS=-g -Wall -lwayland-client `pkg-config --cflags --libs libudev` 
 
 all:
 	gcc main.c ${CFLAGS} -o poll-phone
diff --git a/main.c b/main.c
index 850e6c3..960ee00 100644
--- a/main.c
+++ b/main.c
@@ -10,76 +10,168 @@
 #include <sys/wait.h>
 #include <signal.h>
 #include <errno.h>
+#include <wayland-client.h>
 
 
 #define EPOLL_EVENTS 8
 
+enum app_event {
+    EVENT_UDEV = 0,
+    EVENT_WAYLAND = 1,
+    EVENT_CHILD = 2,
+};
+
+enum app_state {
+    STATE_POLLING,
+    STATE_RUNNING,
+    STATE_FINISHED,
+};
+
 /*Open a connection with the current wayland compositor if it exists
  * A fd is provided which can be polled. When polling the fd returns an error
  * send a signal to any children and terminate.
  */
 
-bool TERMINATE = false;
+struct poll_phone {
+    enum app_state state;
+    pid_t child;
+    int pipe;
+};
+
+void app_abort(char *str) {
+    fprintf(stderr, "%s\n", str); 
+    exit(-1);
+}
+
+static struct poll_phone app;
+
+void run_scrcpy() {
+    assert(app.state != STATE_RUNNING);
 
-int fork_and_wait() {
     pid_t pid = fork();
-    assert(pid != -1);
     if(pid == 0) {
+        signal(SIGTERM, SIG_DFL);
+        signal(SIGINT, SIG_DFL);
         /*There's a race condition with adb...*/
         sleep(2);
-        return execlp("scrcpy", "scrcpy", (char *) 0);
+        execlp("scrcpy", "scrcpy", (char *) 0);
+        exit(-1);
+    } else if(pid == -1) {
+        app_abort("fork() error");
     }
-    int status;
-    pid_t _pid;
-    do {
-        _pid = waitpid(pid, &status, WNOHANG);
-    } while (pid != _pid && !TERMINATE);
-    int exit_status = WEXITSTATUS(status);
-    return exit_status;
+    app.state = STATE_RUNNING;
+    app.child = pid;
 }
 
-void handle_hup(int sig) {
-    printf("poll-phone: %d\n", sig);
+void handle_term(int sig) {
     signal(sig, SIG_IGN);
-    kill(0, SIGHUP);
-    TERMINATE = true;
+    app.state = STATE_FINISHED;
+}
+
+void handle_child(int sig) {
+    char data[1] = {0};
+    app.state = STATE_POLLING;
+    if(write(app.pipe, data, 1) == -1) {
+        app_abort("write failed");
+    }
 }
 
 /*Refactor this so we poll the signal handlers*/
 
 int main(int argc, char **argv) {
 
+    int fileds[2];
+    char buffer[128];
+    if(pipe(fileds) == -1) {
+        app_abort("Failled to create pipe.");
+    }
+    app.state = STATE_POLLING;
+    app.pipe = fileds[1];
+    app.child = -1;
+    signal(SIGTERM, handle_term);
+    signal(SIGINT, handle_term);
+    signal(SIGCHLD, handle_child);
+
     struct epoll_event events[EPOLL_EVENTS];
-    int epfd = epoll_create(EPOLL_EVENTS);
+    struct epoll_event event;
+
+    struct wl_display *display = wl_display_connect(NULL);
+    if(!display) {
+        app_abort("Failled to connect to Wayland display.");
+    }
+
     struct udev *_udev = udev_new();
-    assert(_udev != NULL);
+    if(!_udev) {
+        app_abort("Failed to allocate udev");
+    }
+
     struct udev_monitor *_monitor = udev_monitor_new_from_netlink(_udev, "udev");
-    assert(_monitor != NULL);
+    if(!_monitor) {
+        app_abort("Failed to allocate udev monitor");
+    }
     udev_monitor_enable_receiving(_monitor);
-    int fd = udev_monitor_get_fd(_monitor);
 
-    struct epoll_event event;
-    event.events = EPOLLIN;
-    epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &event);
-
-    signal(SIGTERM, handle_hup);
-    signal(SIGINT, handle_hup);
+    int epfd = epoll_create(EPOLL_EVENTS);
+    if(epfd == -1) {
+        app_abort("Failed to create epoll");
+    }
 
-    fork_and_wait();
-    while(!TERMINATE) {
-        int state = epoll_wait(epfd, &events[0], EPOLL_EVENTS, -1);
-        if(state == EINTR) { break; }
-        else { exit(-1); }
+    event.events = EPOLLIN;
+    event.data.u32 = EVENT_UDEV;
+    epoll_ctl(epfd, EPOLL_CTL_ADD, udev_monitor_get_fd(_monitor), &event);
+    event.data.u32 = EVENT_WAYLAND;
+    epoll_ctl(epfd, EPOLL_CTL_ADD, wl_display_get_fd(display), &event);
+    event.data.u32 = EVENT_CHILD;
+    epoll_ctl(epfd, EPOLL_CTL_ADD, fileds[0], &event);
 
-        struct udev_device *device = udev_monitor_receive_device(_monitor);
-        const char *action = udev_device_get_action(device);
-        udev_device_unref(device);
+    run_scrcpy();
+    while(app.state != STATE_FINISHED) {
+        int ready = epoll_wait(epfd, &events[0], EPOLL_EVENTS, -1);
+        if(ready > 0) {
+            for(int i = 0; i < ready; i++) {
+                struct epoll_event *event = &events[i];
+                switch(event->data.u32) {
+                case EVENT_UDEV: {
+                    struct udev_device *device = udev_monitor_receive_device(_monitor);
+                    const char *action = udev_device_get_action(device);
+                    udev_device_unref(device);
+                    if(!strcmp(action, "bind") && app.state == STATE_POLLING) {
+                        run_scrcpy();
+                    } 
+                    break;
+                    }
+                case EVENT_WAYLAND:
+                    if(wl_display_dispatch(display) == -1) {
+                        app.state = STATE_FINISHED;
+                    }
+                    break;
+                case EVENT_CHILD:
+                    {
+                    if(read(fileds[0], buffer, 128) == -1) {
+                        app_abort("Failed to read");
+                    }
+                    if(app.state != STATE_FINISHED) {
+                        app.state = STATE_POLLING;
+                    }
+                    wait(NULL);
+                    app.child = -1;
+                    }
+                    break;
+                }
+            }
+        }
+    }
 
-        if(!strcmp(action, "bind")) {
-            fork_and_wait();
-        } 
+    kill(0, SIGTERM);
+    if(app.child != -1) {
+        pid_t r;
+        do {
+            r = wait(NULL);
+        } while (r != app.child && r > 0);
     }
 
+    close(fileds[0]);
+    close(fileds[1]);
     udev_monitor_unref(_monitor);
     udev_unref(_udev);
     return 0;