From 36baf441790a12607b61d7b4adeafaa409f3fe6e Mon Sep 17 00:00:00 2001
From: Suren Baghdasaryan <surenb@google.com>
Date: Mon, 23 Dec 2019 11:37:34 -0800
Subject: [PATCH] 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>
---
 include/lmkd.h | 33 ++++++++++++++++++++++++++++++++-
 lmkd.cpp       | 16 +++++++++++++++-
 2 files changed, 47 insertions(+), 2 deletions(-)

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 @@ enum lmk_cmd {
     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 @@ static inline size_t lmkd_pack_set_getkillcnt_repl(LMKD_CTRL_PACKET packet, int
     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 event_handler_info {
 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 @@ static void ctrl_data_write_lmk_kill_occurred(pid_t pid, uid_t uid) {
     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_procpurge(struct ucred *cred) {
     }
 }
 
+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 @@ static void ctrl_command_handler(int dsock_idx) {
         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 @@ static void ctrl_connect_handler(int data __unused, uint32_t events __unused,
     /* 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) {