lmkd: Restrict lmkd unsolicited notifications only to subscribed recipients

lmkd unsolicited notifications can cause lmkd to block if clients are not
consuming them. Fix that by sending notifications to only subscribed
clients. Introduce LMK_SUBSCRIBE command to allow lmkd clients to subscribe
to event notifications. The only asynchronous event currently supported is
LMK_ASYNC_EVENT_KILL.

Bug: 146597855
Test: fill up send buffer using lmkd_unit_test
Test: confirm lmkd does not block after the fix
Change-Id: I014159aa55b59081f4b9ed53ecd160a49c0682bb
Signed-off-by: Suren Baghdasaryan <surenb@google.com>
diff --git a/include/lmkd.h b/include/lmkd.h
index e27494c..c22cb78 100644
--- a/include/lmkd.h
+++ b/include/lmkd.h
@@ -32,7 +32,8 @@
     LMK_PROCREMOVE, /* Unregister a process */
     LMK_PROCPURGE,  /* Purge all registered processes */
     LMK_GETKILLCNT, /* Get number of kills */
-    LMK_PROCKILL,   /* Unsolicited msg to system_server on proc kills */
+    LMK_SUBSCRIBE,  /* Subscribe for asynchronous events */
+    LMK_PROCKILL,   /* Unsolicited msg to subscribed clients on proc kills */
 };
 
 /*
@@ -202,6 +203,36 @@
     return 2 * sizeof(int);
 }
 
+/* Types of asyncronous events sent from lmkd to its clients */
+enum async_event_type {
+    LMK_ASYNC_EVENT_FIRST,
+    LMK_ASYNC_EVENT_KILL = LMK_ASYNC_EVENT_FIRST,
+    LMK_ASYNC_EVENT_COUNT,
+};
+
+/* LMK_SUBSCRIBE packet payload */
+struct lmk_subscribe {
+    enum async_event_type evt_type;
+};
+
+/*
+ * For LMK_SUBSCRIBE packet get its payload.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline void lmkd_pack_get_subscribe(LMKD_CTRL_PACKET packet, struct lmk_subscribe* params) {
+    params->evt_type = (enum async_event_type)ntohl(packet[1]);
+}
+
+/**
+ * Prepare LMK_SUBSCRIBE packet and return packet size in bytes.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline size_t lmkd_pack_set_subscribe(LMKD_CTRL_PACKET packet, enum async_event_type evt_type) {
+    packet[0] = htonl(LMK_SUBSCRIBE);
+    packet[1] = htonl((int)evt_type);
+    return 2 * sizeof(int);
+}
+
 /**
  * Prepare LMK_PROCKILL unsolicited packet and return packet size in bytes.
  * Warning: no checks performed, caller should ensure valid parameters.
diff --git a/lmkd.cpp b/lmkd.cpp
index 389a2de..f7c631d 100644
--- a/lmkd.cpp
+++ b/lmkd.cpp
@@ -242,6 +242,7 @@
 struct sock_event_handler_info {
     int sock;
     pid_t pid;
+    uint32_t async_event_mask;
     struct event_handler_info handler_info;
 };
 
@@ -749,7 +750,7 @@
     size_t len = lmkd_pack_set_prockills(packet, pid, uid);
 
     for (int i = 0; i < MAX_DATA_CONN; i++) {
-        if (data_sock[i].sock >= 0) {
+        if (data_sock[i].sock >= 0 && data_sock[i].async_event_mask & 1 << LMK_ASYNC_EVENT_KILL) {
             ctrl_data_write(i, (char*)packet, len);
         }
     }
@@ -1211,6 +1212,13 @@
     }
 }
 
+static void cmd_subscribe(int dsock_idx, LMKD_CTRL_PACKET packet) {
+    struct lmk_subscribe params;
+
+    lmkd_pack_get_subscribe(packet, &params);
+    data_sock[dsock_idx].async_event_mask |= 1 << params.evt_type;
+}
+
 static void inc_killcnt(int oomadj) {
     int slot = ADJTOSLOT(oomadj);
     uint8_t idx = killcnt_idx[slot];
@@ -1401,6 +1409,11 @@
         if (ctrl_data_write(dsock_idx, (char *)packet, len) != len)
             return;
         break;
+    case LMK_SUBSCRIBE:
+        if (nargs != 1)
+            goto wronglen;
+        cmd_subscribe(dsock_idx, packet);
+        break;
     case LMK_PROCKILL:
         /* This command code is NOT expected at all */
         ALOGE("Received unexpected command code %d", cmd);
@@ -1461,6 +1474,7 @@
     /* use data to store data connection idx */
     data_sock[free_dscock_idx].handler_info.data = free_dscock_idx;
     data_sock[free_dscock_idx].handler_info.handler = ctrl_data_handler;
+    data_sock[free_dscock_idx].async_event_mask = 0;
     epev.events = EPOLLIN;
     epev.data.ptr = (void *)&(data_sock[free_dscock_idx].handler_info);
     if (epoll_ctl(epollfd, EPOLL_CTL_ADD, data_sock[free_dscock_idx].sock, &epev) == -1) {