Merge changes from topic "native process registration"

* changes:
  init: Add support for native service registration with lmkd
  lmkd: Support process types when registering a process
  lmkd: Add library function to unregister a process
  lmkd: Prepare lmkd to support connection from init process
  lmkd: Restrict process record modifications to the client that created it
diff --git a/init/Android.bp b/init/Android.bp
index 776a3a6..d939fcc 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -65,6 +65,7 @@
         "libavb",
         "libc++fs",
         "libcgrouprc_format",
+        "liblmkd_utils",
         "libmodprobe",
         "libprotobuf-cpp-lite",
         "libpropertyinfoserializer",
@@ -118,6 +119,7 @@
         "init.cpp",
         "interface_utils.cpp",
         "keychords.cpp",
+        "lmkd_service.cpp",
         "modalias_handler.cpp",
         "mount_handler.cpp",
         "mount_namespace.cpp",
diff --git a/init/init.cpp b/init/init.cpp
index f775d8f..6ea2d00 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -59,6 +59,7 @@
 #include "first_stage_mount.h"
 #include "import_parser.h"
 #include "keychords.h"
+#include "lmkd_service.h"
 #include "mount_handler.h"
 #include "mount_namespace.h"
 #include "property_service.h"
@@ -684,9 +685,15 @@
     InitKernelLogging(argv);
     LOG(INFO) << "init second stage started!";
 
+    // Will handle EPIPE at the time of write by checking the errno
+    signal(SIGPIPE, SIG_IGN);
+
     // Set init and its forked children's oom_adj.
-    if (auto result = WriteFile("/proc/1/oom_score_adj", "-1000"); !result) {
-        LOG(ERROR) << "Unable to write -1000 to /proc/1/oom_score_adj: " << result.error();
+    if (auto result =
+                WriteFile("/proc/1/oom_score_adj", StringPrintf("%d", DEFAULT_OOM_SCORE_ADJUST));
+        !result) {
+        LOG(ERROR) << "Unable to write " << DEFAULT_OOM_SCORE_ADJUST
+                   << " to /proc/1/oom_score_adj: " << result.error();
     }
 
     // Set up a session keyring that all processes will have access to. It
diff --git a/init/lmkd_service.cpp b/init/lmkd_service.cpp
new file mode 100644
index 0000000..dd1ab4d
--- /dev/null
+++ b/init/lmkd_service.cpp
@@ -0,0 +1,131 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "lmkd_service.h"
+
+#include <errno.h>
+
+#include <android-base/logging.h>
+#include <liblmkd_utils.h>
+
+#include "service_list.h"
+
+namespace android {
+namespace init {
+
+enum LmkdRegistrationResult {
+    LMKD_REG_SUCCESS,
+    LMKD_CONN_FAILED,
+    LMKD_REG_FAILED,
+};
+
+static int lmkd_socket = -1;
+
+static LmkdRegistrationResult RegisterProcess(uid_t uid, pid_t pid, int oom_score_adjust) {
+    // connect to lmkd if not already connected
+    if (lmkd_socket < 0) {
+        lmkd_socket = lmkd_connect();
+        if (lmkd_socket < 0) {
+            return LMKD_CONN_FAILED;
+        }
+    }
+
+    // register service with lmkd
+    struct lmk_procprio params;
+    params.pid = pid;
+    params.uid = uid;
+    params.oomadj = oom_score_adjust;
+    params.ptype = PROC_TYPE_SERVICE;
+    if (lmkd_register_proc(lmkd_socket, &params) != 0) {
+        // data transfer failed, reset the connection
+        close(lmkd_socket);
+        lmkd_socket = -1;
+        return LMKD_REG_FAILED;
+    }
+
+    return LMKD_REG_SUCCESS;
+}
+
+static bool UnregisterProcess(pid_t pid) {
+    if (lmkd_socket < 0) {
+        // no connection or it was lost, no need to unregister
+        return false;
+    }
+
+    // unregister service
+    struct lmk_procremove params;
+    params.pid = pid;
+    if (lmkd_unregister_proc(lmkd_socket, &params) != 0) {
+        // data transfer failed, reset the connection
+        close(lmkd_socket);
+        lmkd_socket = -1;
+        return false;
+    }
+
+    return true;
+}
+
+static void RegisterServices(pid_t exclude_pid) {
+    for (const auto& service : ServiceList::GetInstance().services()) {
+        auto svc = service.get();
+        if (svc->oom_score_adjust() != DEFAULT_OOM_SCORE_ADJUST) {
+            // skip if process is excluded or not yet forked (pid==0)
+            if (svc->pid() == exclude_pid || svc->pid() == 0) {
+                continue;
+            }
+            if (RegisterProcess(svc->uid(), svc->pid(), svc->oom_score_adjust()) !=
+                LMKD_REG_SUCCESS) {
+                // a failure here resets the connection, will retry during next registration
+                break;
+            }
+        }
+    }
+}
+
+void LmkdRegister(const std::string& name, uid_t uid, pid_t pid, int oom_score_adjust) {
+    bool new_connection = lmkd_socket == -1;
+    LmkdRegistrationResult result;
+
+    result = RegisterProcess(uid, pid, oom_score_adjust);
+    if (result == LMKD_REG_FAILED) {
+        // retry one time if connection to lmkd was lost
+        result = RegisterProcess(uid, pid, oom_score_adjust);
+        new_connection = result == LMKD_REG_SUCCESS;
+    }
+    switch (result) {
+        case LMKD_REG_SUCCESS:
+            // register existing services once new connection is established
+            if (new_connection) {
+                RegisterServices(pid);
+            }
+            break;
+        case LMKD_CONN_FAILED:
+            PLOG(ERROR) << "lmkd connection failed when " << name << " process got started";
+            break;
+        case LMKD_REG_FAILED:
+            PLOG(ERROR) << "lmkd failed to register " << name << " process";
+            break;
+    }
+}
+
+void LmkdUnregister(const std::string& name, pid_t pid) {
+    if (!UnregisterProcess(pid)) {
+        PLOG(ERROR) << "lmkd failed to unregister " << name << " process";
+    }
+}
+
+}  // namespace init
+}  // namespace android
diff --git a/init/lmkd_service.h b/init/lmkd_service.h
new file mode 100644
index 0000000..5b51d52
--- /dev/null
+++ b/init/lmkd_service.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <sys/types.h>
+
+#include <string>
+
+namespace android {
+namespace init {
+
+static const int MIN_OOM_SCORE_ADJUST = -1000;
+static const int MAX_OOM_SCORE_ADJUST = 1000;
+// service with default score is unkillable
+static const int DEFAULT_OOM_SCORE_ADJUST = MIN_OOM_SCORE_ADJUST;
+
+#if defined(__ANDROID__)
+
+void LmkdRegister(const std::string& name, uid_t uid, pid_t pid, int oom_score_adjust);
+void LmkdUnregister(const std::string& name, pid_t pid);
+
+#else  // defined(__ANDROID__)
+
+static inline void LmkdRegister(const std::string&, uid_t, pid_t, int) {}
+static inline void LmkdUnregister(const std::string&, pid_t) {}
+
+#endif  // defined(__ANDROID__)
+
+}  // namespace init
+}  // namespace android
diff --git a/init/service.cpp b/init/service.cpp
index c8568a0..f8e98a2 100644
--- a/init/service.cpp
+++ b/init/service.cpp
@@ -36,6 +36,7 @@
 #include <processgroup/processgroup.h>
 #include <selinux/selinux.h>
 
+#include "lmkd_service.h"
 #include "service_list.h"
 #include "util.h"
 
@@ -151,7 +152,7 @@
       seclabel_(seclabel),
       onrestart_(false, subcontext_for_restart_commands, "<Service '" + name + "' onrestart>", 0,
                  "onrestart", {}),
-      oom_score_adjust_(-1000),
+      oom_score_adjust_(DEFAULT_OOM_SCORE_ADJUST),
       start_order_(0),
       args_(args) {}
 
@@ -199,6 +200,10 @@
 
         if (r == 0) process_cgroup_empty_ = true;
     }
+
+    if (oom_score_adjust_ != DEFAULT_OOM_SCORE_ADJUST) {
+        LmkdUnregister(name_, pid_);
+    }
 }
 
 void Service::SetProcessAttributesAndCaps() {
@@ -502,7 +507,7 @@
         return ErrnoError() << "Failed to fork";
     }
 
-    if (oom_score_adjust_ != -1000) {
+    if (oom_score_adjust_ != DEFAULT_OOM_SCORE_ADJUST) {
         std::string oom_str = std::to_string(oom_score_adjust_);
         std::string oom_file = StringPrintf("/proc/%d/oom_score_adj", pid);
         if (!WriteStringToFile(oom_str, oom_file)) {
@@ -563,6 +568,10 @@
         }
     }
 
+    if (oom_score_adjust_ != DEFAULT_OOM_SCORE_ADJUST) {
+        LmkdRegister(name_, proc_attr_.uid, pid_, oom_score_adjust_);
+    }
+
     NotifyStateChange("running");
     reboot_on_failure.Disable();
     return {};
diff --git a/init/service_parser.cpp b/init/service_parser.cpp
index e6a341d..154d1dd 100644
--- a/init/service_parser.cpp
+++ b/init/service_parser.cpp
@@ -29,6 +29,7 @@
 #include <hidl-util/FQName.h>
 #include <system/thread_defs.h>
 
+#include "lmkd_service.h"
 #include "rlimit_parser.h"
 #include "service_utils.h"
 #include "util.h"
@@ -261,8 +262,10 @@
 }
 
 Result<void> ServiceParser::ParseOomScoreAdjust(std::vector<std::string>&& args) {
-    if (!ParseInt(args[1], &service_->oom_score_adjust_, -1000, 1000)) {
-        return Error() << "oom_score_adjust value must be in range -1000 - +1000";
+    if (!ParseInt(args[1], &service_->oom_score_adjust_, MIN_OOM_SCORE_ADJUST,
+                  MAX_OOM_SCORE_ADJUST)) {
+        return Error() << "oom_score_adjust value must be in range " << MIN_OOM_SCORE_ADJUST
+                       << " - +" << MAX_OOM_SCORE_ADJUST;
     }
     return {};
 }
diff --git a/init/service_test.cpp b/init/service_test.cpp
index c9cc7bd..c158b0a 100644
--- a/init/service_test.cpp
+++ b/init/service_test.cpp
@@ -23,6 +23,7 @@
 
 #include <gtest/gtest.h>
 
+#include "lmkd_service.h"
 #include "util.h"
 
 namespace android {
@@ -49,7 +50,7 @@
     EXPECT_EQ(IoSchedClass_NONE, service_in_old_memory->ioprio_class());
     EXPECT_EQ(0, service_in_old_memory->ioprio_pri());
     EXPECT_EQ(0, service_in_old_memory->priority());
-    EXPECT_EQ(-1000, service_in_old_memory->oom_score_adjust());
+    EXPECT_EQ(DEFAULT_OOM_SCORE_ADJUST, service_in_old_memory->oom_score_adjust());
     EXPECT_FALSE(service_in_old_memory->process_cgroup_empty());
 
     for (std::size_t i = 0; i < memory_size; ++i) {
@@ -68,7 +69,7 @@
     EXPECT_EQ(IoSchedClass_NONE, service_in_old_memory2->ioprio_class());
     EXPECT_EQ(0, service_in_old_memory2->ioprio_pri());
     EXPECT_EQ(0, service_in_old_memory2->priority());
-    EXPECT_EQ(-1000, service_in_old_memory2->oom_score_adjust());
+    EXPECT_EQ(DEFAULT_OOM_SCORE_ADJUST, service_in_old_memory2->oom_score_adjust());
     EXPECT_FALSE(service_in_old_memory->process_cgroup_empty());
 }
 
diff --git a/lmkd/Android.bp b/lmkd/Android.bp
index e8e125b..deebb8e 100644
--- a/lmkd/Android.bp
+++ b/lmkd/Android.bp
@@ -48,6 +48,7 @@
 cc_library_static {
     name: "liblmkd_utils",
     srcs: ["liblmkd_utils.c"],
+    recovery_available: true,
     shared_libs: [
         "libcutils",
     ],
diff --git a/lmkd/include/liblmkd_utils.h b/lmkd/include/liblmkd_utils.h
index 72e3f4a..92e4d41 100644
--- a/lmkd/include/liblmkd_utils.h
+++ b/lmkd/include/liblmkd_utils.h
@@ -40,6 +40,14 @@
 int lmkd_register_proc(int sock, struct lmk_procprio *params);
 
 /*
+ * Unregisters a process previously registered with lmkd.
+ * On success returns 0.
+ * On error, -1 is returned.
+ * In the case of error errno is set appropriately.
+ */
+int lmkd_unregister_proc(int sock, struct lmk_procremove *params);
+
+/*
  * Creates memcg directory for given process.
  * On success returns 0.
  * -1 is returned if path creation failed.
diff --git a/lmkd/include/lmkd.h b/lmkd/include/lmkd.h
index 59377dd..bd9b80e 100644
--- a/lmkd/include/lmkd.h
+++ b/lmkd/include/lmkd.h
@@ -87,21 +87,33 @@
     return idx * sizeof(int);
 }
 
+/* Process types for lmk_procprio.ptype */
+enum proc_type {
+    PROC_TYPE_FIRST,
+    PROC_TYPE_APP = PROC_TYPE_FIRST,
+    PROC_TYPE_SERVICE,
+    PROC_TYPE_COUNT,
+};
+
 /* LMK_PROCPRIO packet payload */
 struct lmk_procprio {
     pid_t pid;
     uid_t uid;
     int oomadj;
+    enum proc_type ptype;
 };
 
 /*
  * For LMK_PROCPRIO packet get its payload.
  * Warning: no checks performed, caller should ensure valid parameters.
  */
-static inline void lmkd_pack_get_procprio(LMKD_CTRL_PACKET packet, struct lmk_procprio* params) {
+static inline void lmkd_pack_get_procprio(LMKD_CTRL_PACKET packet, int field_count,
+                                          struct lmk_procprio* params) {
     params->pid = (pid_t)ntohl(packet[1]);
     params->uid = (uid_t)ntohl(packet[2]);
     params->oomadj = ntohl(packet[3]);
+    /* if field is missing assume PROC_TYPE_APP for backward compatibility */
+    params->ptype = field_count > 3 ? (enum proc_type)ntohl(packet[4]) : PROC_TYPE_APP;
 }
 
 /*
@@ -113,7 +125,8 @@
     packet[1] = htonl(params->pid);
     packet[2] = htonl(params->uid);
     packet[3] = htonl(params->oomadj);
-    return 4 * sizeof(int);
+    packet[4] = htonl((int)params->ptype);
+    return 5 * sizeof(int);
 }
 
 /* LMK_PROCREMOVE packet payload */
@@ -135,7 +148,7 @@
  * Warning: no checks performed, caller should ensure valid parameters.
  */
 static inline size_t lmkd_pack_set_procremove(LMKD_CTRL_PACKET packet,
-                                              struct lmk_procprio* params) {
+                                              struct lmk_procremove* params) {
     packet[0] = htonl(LMK_PROCREMOVE);
     packet[1] = htonl(params->pid);
     return 2 * sizeof(int);
diff --git a/lmkd/liblmkd_utils.c b/lmkd/liblmkd_utils.c
index fa3b7a9..280c149 100644
--- a/lmkd/liblmkd_utils.c
+++ b/lmkd/liblmkd_utils.c
@@ -28,7 +28,7 @@
 int lmkd_connect() {
     return socket_local_client("lmkd",
                                ANDROID_SOCKET_NAMESPACE_RESERVED,
-                               SOCK_SEQPACKET);
+                               SOCK_SEQPACKET | SOCK_CLOEXEC);
 }
 
 int lmkd_register_proc(int sock, struct lmk_procprio *params) {
@@ -42,6 +42,17 @@
     return (ret < 0) ? -1 : 0;
 }
 
+int lmkd_unregister_proc(int sock, struct lmk_procremove *params) {
+    LMKD_CTRL_PACKET packet;
+    size_t size;
+    int ret;
+
+    size = lmkd_pack_set_procremove(packet, params);
+    ret = TEMP_FAILURE_RETRY(write(sock, packet, size));
+
+    return (ret < 0) ? -1 : 0;
+}
+
 int create_memcg(uid_t uid, pid_t pid) {
     char buf[256];
     int tasks_file;
diff --git a/lmkd/lmkd.c b/lmkd/lmkd.c
index 9de7ff7..4352498 100644
--- a/lmkd/lmkd.c
+++ b/lmkd/lmkd.c
@@ -45,6 +45,7 @@
 #include <log/log.h>
 #include <log/log_event_list.h>
 #include <log/log_time.h>
+#include <private/android_filesystem_config.h>
 #include <psi/psi.h>
 #include <system/thread_defs.h>
 
@@ -240,11 +241,12 @@
 /* data required to handle socket events */
 struct sock_event_handler_info {
     int sock;
+    pid_t pid;
     struct event_handler_info handler_info;
 };
 
-/* max supported number of data connections */
-#define MAX_DATA_CONN 2
+/* max supported number of data connections (AMS, init, tests) */
+#define MAX_DATA_CONN 3
 
 /* socket event handler data */
 static struct sock_event_handler_info ctrl_sock;
@@ -254,7 +256,7 @@
 static struct event_handler_info vmpressure_hinfo[VMPRESS_LEVEL_COUNT];
 
 /*
- * 1 ctrl listen socket, 2 ctrl data socket, 3 memory pressure levels,
+ * 1 ctrl listen socket, 3 ctrl data socket, 3 memory pressure levels,
  * 1 lmk events + 1 fd to wait for process death
  */
 #define MAX_EPOLL_EVENTS (1 + MAX_DATA_CONN + VMPRESS_LEVEL_COUNT + 1 + 1)
@@ -490,6 +492,7 @@
     int pidfd;
     uid_t uid;
     int oomadj;
+    pid_t reg_pid; /* PID of the process that registered this record */
     struct proc *pidhash_next;
 };
 
@@ -845,7 +848,35 @@
     return buf;
 }
 
-static void cmd_procprio(LMKD_CTRL_PACKET packet) {
+static bool claim_record(struct proc *procp, pid_t pid) {
+    if (procp->reg_pid == pid) {
+        /* Record already belongs to the registrant */
+        return true;
+    }
+    if (procp->reg_pid == 0) {
+        /* Old registrant is gone, claim the record */
+        procp->reg_pid = pid;
+        return true;
+    }
+    /* The record is owned by another registrant */
+    return false;
+}
+
+static void remove_claims(pid_t pid) {
+    int i;
+
+    for (i = 0; i < PIDHASH_SZ; i++) {
+        struct proc *procp = pidhash[i];
+        while (procp) {
+            if (procp->reg_pid == pid) {
+                procp->reg_pid = 0;
+            }
+            procp = procp->pidhash_next;
+        }
+    }
+}
+
+static void cmd_procprio(LMKD_CTRL_PACKET packet, int field_count, struct ucred *cred) {
     struct proc *procp;
     char path[LINE_MAX];
     char val[20];
@@ -855,7 +886,7 @@
     struct passwd *pwdrec;
     int tgid;
 
-    lmkd_pack_get_procprio(packet, &params);
+    lmkd_pack_get_procprio(packet, field_count, &params);
 
     if (params.oomadj < OOM_SCORE_ADJ_MIN ||
         params.oomadj > OOM_SCORE_ADJ_MAX) {
@@ -863,6 +894,11 @@
         return;
     }
 
+    if (params.ptype < PROC_TYPE_FIRST || params.ptype >= PROC_TYPE_COUNT) {
+        ALOGE("Invalid PROCPRIO process type argument %d", params.ptype);
+        return;
+    }
+
     /* Check if registered process is a thread group leader */
     tgid = proc_get_tgid(params.pid);
     if (tgid >= 0 && tgid != params.pid) {
@@ -889,7 +925,8 @@
         return;
     }
 
-    if (per_app_memcg) {
+    /* lmkd should not change soft limits for services */
+    if (params.ptype == PROC_TYPE_APP && per_app_memcg) {
         if (params.oomadj >= 900) {
             soft_limit_mult = 0;
         } else if (params.oomadj >= 800) {
@@ -954,24 +991,47 @@
         procp->pid = params.pid;
         procp->pidfd = pidfd;
         procp->uid = params.uid;
+        procp->reg_pid = cred->pid;
         procp->oomadj = params.oomadj;
         proc_insert(procp);
     } else {
+        if (!claim_record(procp, cred->pid)) {
+            char buf[LINE_MAX];
+            /* Only registrant of the record can remove it */
+            ALOGE("%s (%d, %d) attempts to modify a process registered by another client",
+                proc_get_name(cred->pid, buf, sizeof(buf)), cred->uid, cred->pid);
+            return;
+        }
         proc_unslot(procp);
         procp->oomadj = params.oomadj;
         proc_slot(procp);
     }
 }
 
-static void cmd_procremove(LMKD_CTRL_PACKET packet) {
+static void cmd_procremove(LMKD_CTRL_PACKET packet, struct ucred *cred) {
     struct lmk_procremove params;
+    struct proc *procp;
 
     lmkd_pack_get_procremove(packet, &params);
+
     if (use_inkernel_interface) {
         stats_remove_taskname(params.pid, kpoll_info.poll_fd);
         return;
     }
 
+    procp = pid_lookup(params.pid);
+    if (!procp) {
+        return;
+    }
+
+    if (!claim_record(procp, cred->pid)) {
+        char buf[LINE_MAX];
+        /* Only registrant of the record can remove it */
+        ALOGE("%s (%d, %d) attempts to unregister a process registered by another client",
+            proc_get_name(cred->pid, buf, sizeof(buf)), cred->uid, cred->pid);
+        return;
+    }
+
     /*
      * WARNING: After pid_remove() procp is freed and can't be used!
      * Therefore placed at the end of the function.
@@ -979,7 +1039,7 @@
     pid_remove(params.pid);
 }
 
-static void cmd_procpurge() {
+static void cmd_procpurge(struct ucred *cred) {
     int i;
     struct proc *procp;
     struct proc *next;
@@ -989,20 +1049,17 @@
         return;
     }
 
-    for (i = 0; i <= ADJTOSLOT(OOM_SCORE_ADJ_MAX); i++) {
-        procadjslot_list[i].next = &procadjslot_list[i];
-        procadjslot_list[i].prev = &procadjslot_list[i];
-    }
-
     for (i = 0; i < PIDHASH_SZ; i++) {
         procp = pidhash[i];
         while (procp) {
             next = procp->pidhash_next;
-            free(procp);
+            /* Purge only records created by the requestor */
+            if (claim_record(procp, cred->pid)) {
+                pid_remove(procp->pid);
+            }
             procp = next;
         }
     }
-    memset(&pidhash[0], 0, sizeof(pidhash));
 }
 
 static void inc_killcnt(int oomadj) {
@@ -1153,19 +1210,50 @@
 
     close(data_sock[dsock_idx].sock);
     data_sock[dsock_idx].sock = -1;
+
+    /* Mark all records of the old registrant as unclaimed */
+    remove_claims(data_sock[dsock_idx].pid);
 }
 
-static int ctrl_data_read(int dsock_idx, char *buf, size_t bufsz) {
-    int ret = 0;
+static ssize_t ctrl_data_read(int dsock_idx, char *buf, size_t bufsz, struct ucred *sender_cred) {
+    struct iovec iov = { buf, bufsz };
+    char control[CMSG_SPACE(sizeof(struct ucred))];
+    struct msghdr hdr = {
+        NULL, 0, &iov, 1, control, sizeof(control), 0,
+    };
+    ssize_t ret;
 
-    ret = TEMP_FAILURE_RETRY(read(data_sock[dsock_idx].sock, buf, bufsz));
-
+    ret = TEMP_FAILURE_RETRY(recvmsg(data_sock[dsock_idx].sock, &hdr, 0));
     if (ret == -1) {
-        ALOGE("control data socket read failed; errno=%d", errno);
-    } else if (ret == 0) {
-        ALOGE("Got EOF on control data socket");
-        ret = -1;
+        ALOGE("control data socket read failed; %s", strerror(errno));
+        return -1;
     }
+    if (ret == 0) {
+        ALOGE("Got EOF on control data socket");
+        return -1;
+    }
+
+    struct ucred* cred = NULL;
+    struct cmsghdr* cmsg = CMSG_FIRSTHDR(&hdr);
+    while (cmsg != NULL) {
+        if (cmsg->cmsg_level == SOL_SOCKET && cmsg->cmsg_type == SCM_CREDENTIALS) {
+            cred = (struct ucred*)CMSG_DATA(cmsg);
+            break;
+        }
+        cmsg = CMSG_NXTHDR(&hdr, cmsg);
+    }
+
+    if (cred == NULL) {
+        ALOGE("Failed to retrieve sender credentials");
+        /* Close the connection */
+        ctrl_data_close(dsock_idx);
+        return -1;
+    }
+
+    memcpy(sender_cred, cred, sizeof(struct ucred));
+
+    /* Store PID of the peer */
+    data_sock[dsock_idx].pid = cred->pid;
 
     return ret;
 }
@@ -1187,13 +1275,14 @@
 
 static void ctrl_command_handler(int dsock_idx) {
     LMKD_CTRL_PACKET packet;
+    struct ucred cred;
     int len;
     enum lmk_cmd cmd;
     int nargs;
     int targets;
     int kill_cnt;
 
-    len = ctrl_data_read(dsock_idx, (char *)packet, CTRL_PACKET_MAX_SIZE);
+    len = ctrl_data_read(dsock_idx, (char *)packet, CTRL_PACKET_MAX_SIZE, &cred);
     if (len <= 0)
         return;
 
@@ -1215,19 +1304,20 @@
         cmd_target(targets, packet);
         break;
     case LMK_PROCPRIO:
-        if (nargs != 3)
+        /* process type field is optional for backward compatibility */
+        if (nargs < 3 || nargs > 4)
             goto wronglen;
-        cmd_procprio(packet);
+        cmd_procprio(packet, nargs, &cred);
         break;
     case LMK_PROCREMOVE:
         if (nargs != 1)
             goto wronglen;
-        cmd_procremove(packet);
+        cmd_procremove(packet, &cred);
         break;
     case LMK_PROCPURGE:
         if (nargs != 0)
             goto wronglen;
-        cmd_procpurge();
+        cmd_procpurge(&cred);
         break;
     case LMK_GETKILLCNT:
         if (nargs != 2)
diff --git a/lmkd/lmkd.rc b/lmkd/lmkd.rc
index 76b6055..982a188 100644
--- a/lmkd/lmkd.rc
+++ b/lmkd/lmkd.rc
@@ -4,5 +4,5 @@
     group lmkd system readproc
     capabilities DAC_OVERRIDE KILL IPC_LOCK SYS_NICE SYS_RESOURCE
     critical
-    socket lmkd seqpacket 0660 system system
+    socket lmkd seqpacket+passcred 0660 system system
     writepid /dev/cpuset/system-background/tasks
diff --git a/rootdir/init.rc b/rootdir/init.rc
index e037a43..0827247 100644
--- a/rootdir/init.rc
+++ b/rootdir/init.rc
@@ -347,6 +347,8 @@
 
     # Start logd before any other services run to ensure we capture all of their logs.
     start logd
+    # Start lmkd before any other services run so that it can register them
+    start lmkd
 
     # Start essential services.
     start servicemanager