Add more tests to libdmabufinfo

Add some test to verify the refcount and fd reference is correct when
the dma_buf is shared between processes.

Bug; 63860998
Test: libdmabufinfo_test

Change-Id: Id22e68e7a65820f19847b2faab11c78e6d942d92

libmeminfo/libdmabufinfo/dmabufinfo_test.cpp

@@ -14,14 +14,17 @@
  */
 
 #include <gtest/gtest.h>
+#include <inttypes.h>
 #include <linux/dma-buf.h>
+#include <poll.h>
 #include <sys/mman.h>
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <fstream>
 #include <string>
-#include <vector>
 #include <unordered_map>
+#include <vector>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
@@ -50,6 +53,115 @@ struct ion_heap_data {
 #define DMA_BUF_SET_NAME _IOW(DMA_BUF_BASE, 5, const char*)
 #endif
 
+class fd_sharer {
+  public:
+    fd_sharer();
+    ~fd_sharer() { kill(); }
+
+    bool ok() const { return child_pid > 0; }
+    bool sendfd(int fd);
+    bool kill();
+    pid_t pid() const { return child_pid; }
+
+  private:
+    unique_fd parent_fd, child_fd;
+    pid_t child_pid;
+
+    void run();
+};
+
+fd_sharer::fd_sharer() : parent_fd{}, child_fd{}, child_pid{-1} {
+    bool sp_ok = android::base::Socketpair(SOCK_STREAM, &parent_fd, &child_fd);
+    if (!sp_ok) return;
+
+    child_pid = fork();
+    if (child_pid < 0) return;
+
+    if (child_pid == 0) run();
+}
+
+bool fd_sharer::kill() {
+    int err = ::kill(child_pid, SIGKILL);
+    if (err < 0) return false;
+
+    return ::waitpid(child_pid, nullptr, 0) == child_pid;
+}
+
+void fd_sharer::run() {
+    while (true) {
+        int fd;
+        char unused = 0;
+
+        iovec iov{};
+        iov.iov_base = &unused;
+        iov.iov_len = sizeof(unused);
+
+        msghdr msg{};
+        msg.msg_iov = &iov;
+        msg.msg_iovlen = 1;
+
+        char cmsg_buf[CMSG_SPACE(sizeof(fd))];
+        msg.msg_control = cmsg_buf;
+        msg.msg_controllen = sizeof(cmsg_buf);
+
+        cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
+        cmsg->cmsg_level = SOL_SOCKET;
+        cmsg->cmsg_type = SCM_RIGHTS;
+        cmsg->cmsg_len = CMSG_LEN(sizeof(fd));
+
+        ssize_t s = TEMP_FAILURE_RETRY(recvmsg(child_fd, &msg, 0));
+        if (s == -1) break;
+
+        s = TEMP_FAILURE_RETRY(write(child_fd, &unused, sizeof(unused)));
+        if (s == -1) break;
+    }
+}
+
+bool fd_sharer::sendfd(int fd) {
+    char unused = 0;
+
+    iovec iov{};
+    iov.iov_base = &unused;
+    iov.iov_len = sizeof(unused);
+
+    msghdr msg{};
+    msg.msg_iov = &iov;
+    msg.msg_iovlen = 1;
+
+    char cmsg_buf[CMSG_SPACE(sizeof(fd))];
+    msg.msg_control = cmsg_buf;
+    msg.msg_controllen = sizeof(cmsg_buf);
+
+    cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
+    cmsg->cmsg_level = SOL_SOCKET;
+    cmsg->cmsg_type = SCM_RIGHTS;
+    cmsg->cmsg_len = CMSG_LEN(sizeof(fd));
+
+    int* fd_buf = reinterpret_cast<int*>(CMSG_DATA(cmsg));
+    *fd_buf = fd;
+
+    ssize_t s = TEMP_FAILURE_RETRY(sendmsg(parent_fd, &msg, 0));
+    if (s == -1) return false;
+
+    // The target process installs the fd into its fd table during recvmsg().
+    // So if we return now, there's a brief window between sendfd() finishing
+    // and libmemoryinfo actually seeing that the buffer has been shared.  This
+    // window is just large enough to break tests.
+    //
+    // To work around this, wait for the target process to respond with a dummy
+    // byte, with a timeout of 1 s.
+    pollfd p{};
+    p.fd = parent_fd;
+    p.events = POLL_IN;
+    int ready = poll(&p, 1, 1000);
+    if (ready != 1) return false;
+
+    s = TEMP_FAILURE_RETRY(read(parent_fd, &unused, sizeof(unused)));
+    if (s == -1) return false;
+
+    return true;
+}
+
 #define EXPECT_ONE_BUF_EQ(_bufptr, _name, _fdrefs, _maprefs, _expname, _count, _size) \
     do {                                                                              \
         EXPECT_EQ(_bufptr->name(), _name);                                            \
@@ -140,6 +252,24 @@ class DmaBufTester : public ::testing::Test {
         return unique_fd{fd};
     }
 
+    void readAndCheckDmaBuffer(std::vector<DmaBuffer>* dmabufs, pid_t pid, const std::string name,
+                               size_t fdrefs_size, size_t maprefs_size, const std::string exporter,
+                               size_t refcount, uint64_t buf_size, bool expectFdrefs,
+                               bool expectMapRefs) {
+        EXPECT_TRUE(ReadDmaBufInfo(pid, dmabufs));
+        EXPECT_EQ(dmabufs->size(), 1UL);
+        EXPECT_ONE_BUF_EQ(dmabufs->begin(), name, fdrefs_size, maprefs_size, exporter, refcount,
+                          buf_size);
+        // Make sure the buffer has the right pid too.
+        EXPECT_PID_IN_FDREFS(dmabufs->begin(), pid, expectFdrefs);
+        EXPECT_PID_IN_MAPREFS(dmabufs->begin(), pid, expectMapRefs);
+    }
+
+    bool checkPidRef(DmaBuffer& dmabuf, pid_t pid, int expectFdrefs) {
+        int fdrefs = dmabuf.fdrefs().find(pid)->second;
+        return fdrefs == expectFdrefs;
+    }
+
   private:
     int get_ion_heap_mask() {
         if (ion_fd < 0) {
@@ -196,7 +326,7 @@ TEST_F(DmaBufTester, TestFdRef) {
         EXPECT_ONE_BUF_EQ(dmabufs.begin(), "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL);
 
         // Make sure the buffer has the right pid too.
-        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, false);
+        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, true);
     }
 
     // Now make sure the buffer has disappeared
@@ -222,8 +352,8 @@ TEST_F(DmaBufTester, TestMapRef) {
         EXPECT_ONE_BUF_EQ(dmabufs.begin(), "dmabuftester-4k", 1UL, 1UL, "ion", 2UL, 4096ULL);
 
         // Make sure the buffer has the right pid too.
-        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, false);
-        EXPECT_PID_IN_MAPREFS(dmabufs.begin(), pid, false);
+        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, true);
+        EXPECT_PID_IN_MAPREFS(dmabufs.begin(), pid, true);
 
         // close the file descriptor and re-read the stats
         buf.reset(-1);
@@ -232,8 +362,8 @@ TEST_F(DmaBufTester, TestMapRef) {
         EXPECT_EQ(dmabufs.size(), 1UL);
         EXPECT_ONE_BUF_EQ(dmabufs.begin(), "<unknown>", 0UL, 1UL, "<unknown>", 0UL, 4096ULL);
 
-        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, true);
-        EXPECT_PID_IN_MAPREFS(dmabufs.begin(), pid, false);
+        EXPECT_PID_IN_FDREFS(dmabufs.begin(), pid, false);
+        EXPECT_PID_IN_MAPREFS(dmabufs.begin(), pid, true);
 
         // unmap the bufer and lose all references
         munmap(ptr, 4096);
@@ -244,6 +374,109 @@ TEST_F(DmaBufTester, TestMapRef) {
     EXPECT_TRUE(dmabufs.empty());
 }
 
+TEST_F(DmaBufTester, TestSharedfd) {
+    // Each time a shared buffer is received over a socket, the remote process
+    // will take an extra reference on it.
+
+    ASSERT_TRUE(is_valid());
+
+    pid_t pid = getpid();
+    std::vector<DmaBuffer> dmabufs;
+    {
+        fd_sharer sharer{};
+        ASSERT_TRUE(sharer.ok());
+        // Allocate one buffer and make sure the library can see it
+        unique_fd buf = allocate(4096, "dmabuftester-4k");
+        ASSERT_GT(buf, 0) << "Allocated buffer is invalid";
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+
+        ASSERT_TRUE(sharer.sendfd(buf));
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 2UL, 4096ULL, true,
+                              false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], pid, 1));
+        readAndCheckDmaBuffer(&dmabufs, sharer.pid(), "dmabuftester-4k", 1UL, 0UL, "ion", 2UL,
+                              4096ULL, true, false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], sharer.pid(), 1));
+
+        ASSERT_TRUE(sharer.sendfd(buf));
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 3UL, 4096ULL, true,
+                              false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], pid, 1));
+        readAndCheckDmaBuffer(&dmabufs, sharer.pid(), "dmabuftester-4k", 1UL, 0UL, "ion", 3UL,
+                              4096ULL, true, false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], sharer.pid(), 2));
+
+        ASSERT_TRUE(sharer.kill());
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+    }
+
+    // Now make sure the buffer has disappeared
+    ASSERT_TRUE(ReadDmaBufInfo(pid, &dmabufs));
+    EXPECT_TRUE(dmabufs.empty());
+}
+
+TEST_F(DmaBufTester, DupFdTest) {
+    // dup()ing an fd will make this process take an extra reference on the
+    // shared buffer.
+
+    ASSERT_TRUE(is_valid());
+
+    pid_t pid = getpid();
+    std::vector<DmaBuffer> dmabufs;
+    {
+        // Allocate one buffer and make sure the library can see it
+        unique_fd buf = allocate(4096, "dmabuftester-4k");
+        ASSERT_GT(buf, 0) << "Allocated buffer is invalid";
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+
+        unique_fd buf2{dup(buf)};
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 2UL, 4096ULL, true,
+                              false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], pid, 2));
+
+        close(buf2.release());
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+        EXPECT_TRUE(checkPidRef(dmabufs[0], pid, 1));
+    }
+
+    // Now make sure the buffer has disappeared
+    ASSERT_TRUE(ReadDmaBufInfo(pid, &dmabufs));
+    EXPECT_TRUE(dmabufs.empty());
+}
+
+TEST_F(DmaBufTester, ForkTest) {
+    // fork()ing a child will cause the child to automatically take a reference
+    // on any existing shared buffers.
+    ASSERT_TRUE(is_valid());
+
+    pid_t pid = getpid();
+    std::vector<DmaBuffer> dmabufs;
+    {
+        // Allocate one buffer and make sure the library can see it
+        unique_fd buf = allocate(4096, "dmabuftester-4k");
+        ASSERT_GT(buf, 0) << "Allocated buffer is invalid";
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+        fd_sharer sharer{};
+        ASSERT_TRUE(sharer.ok());
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 2UL, 4096ULL, true,
+                              false);
+        readAndCheckDmaBuffer(&dmabufs, sharer.pid(), "dmabuftester-4k", 1UL, 0UL, "ion", 2UL,
+                              4096ULL, true, false);
+        ASSERT_TRUE(sharer.kill());
+        readAndCheckDmaBuffer(&dmabufs, pid, "dmabuftester-4k", 1UL, 0UL, "ion", 1UL, 4096ULL, true,
+                              false);
+    }
+
+    // Now make sure the buffer has disappeared
+    ASSERT_TRUE(ReadDmaBufInfo(pid, &dmabufs));
+    EXPECT_TRUE(dmabufs.empty());
+}
+
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
     ::android::base::InitLogging(argv, android::base::StderrLogger);

libmeminfo/libdmabufinfo/include/dmabufinfo/dmabufinfo.h

@@ -44,7 +44,7 @@ struct DmaBuffer {
     }
 
     // Getters for each property
-    uint64_t size() { return size_; }
+    uint64_t size() const { return size_; }
     const std::unordered_map<pid_t, int>& fdrefs() const { return fdrefs_; }
     const std::unordered_map<pid_t, int>& maprefs() const { return maprefs_; }
     ino_t inode() const { return inode_; }
@@ -56,6 +56,11 @@ struct DmaBuffer {
     void SetExporter(const std::string& exporter) { exporter_ = exporter; }
     void SetCount(uint64_t count) { count_ = count; }
 
+    bool operator==(const DmaBuffer& rhs) {
+        return (inode_ == rhs.inode()) && (size_ == rhs.size()) && (name_ == rhs.name()) &&
+               (exporter_ == rhs.exporter());
+    }
+
   private:
     ino_t inode_;
     uint64_t size_;