From 68abd1cb29095b3384c33e8d676ffa252d8f9153 Mon Sep 17 00:00:00 2001
From: Tom <tomut@yahoo.com>
Date: Tue, 24 Nov 2020 10:26:32 -0700
Subject: [PATCH] Kernel: Fix SharedBuffer reference counting on fork

We need to not only add a record for a reference, but we need
to copy the reference count on fork as well, because the code
in the fork assumes that it has the same amount of references,
still.

Also, once all references are dropped when a process is disowned,
delete the shared buffer.

Fixes #4076
---
 Kernel/SharedBuffer.cpp   | 18 +++++++++++++-----
 Kernel/SharedBuffer.h     |  5 +++--
 Kernel/Syscalls/shbuf.cpp |  8 ++++++--
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/Kernel/SharedBuffer.cpp b/Kernel/SharedBuffer.cpp
index 39cd8681aa..0e0682a590 100644
--- a/Kernel/SharedBuffer.cpp
+++ b/Kernel/SharedBuffer.cpp
@@ -126,11 +126,14 @@ void SharedBuffer::share_all_shared_buffers(Process& from_process, Process& with
     LOCKER(shared_buffers().lock());
     for (auto& shbuf : shared_buffers().resource()) {
         auto& shared_buffer = *shbuf.value;
-        if (shared_buffer.m_global)
-            continue;
+        // We need to clone all references (including for global shared buffers),
+        // and the reference counts as well.
         for (auto& ref : shared_buffer.m_refs) {
             if (ref.pid == from_process.pid()) {
-                shared_buffer.share_with(with_process.pid());
+                auto ref_count = ref.count;
+                shared_buffer.m_refs.append(Reference(with_process.pid(), ref_count));
+                // NOTE: ref may become invalid after we appended!
+                shared_buffer.m_total_refs += ref_count;
                 break;
             }
         }
@@ -143,7 +146,9 @@ void SharedBuffer::deref_for_process(Process& process)
     for (size_t i = 0; i < m_refs.size(); ++i) {
         auto& ref = m_refs[i];
         if (ref.pid == process.pid()) {
+            ASSERT(ref.count > 0);
             ref.count--;
+            ASSERT(m_total_refs > 0);
             m_total_refs--;
             if (ref.count == 0) {
 #ifdef SHARED_BUFFER_DEBUG
@@ -164,7 +169,7 @@ void SharedBuffer::deref_for_process(Process& process)
     ASSERT_NOT_REACHED();
 }
 
-void SharedBuffer::disown(ProcessID pid)
+bool SharedBuffer::disown(ProcessID pid)
 {
     LOCKER(shared_buffers().lock());
     for (size_t i = 0; i < m_refs.size(); ++i) {
@@ -173,15 +178,18 @@ void SharedBuffer::disown(ProcessID pid)
 #ifdef SHARED_BUFFER_DEBUG
             dbg() << "Disowning shared buffer " << m_shbuf_id << " of size " << size() << " by PID " << pid.value();
 #endif
+            ASSERT(m_total_refs >= ref.count);
             m_total_refs -= ref.count;
             m_refs.unstable_take(i);
 #ifdef SHARED_BUFFER_DEBUG
             dbg() << "Disowned shared buffer " << m_shbuf_id << " of size " << size() << " by PID " << pid.value();
 #endif
             destroy_if_unused();
-            return;
+            break;
         }
     }
+
+    return m_total_refs == 0;
 }
 
 void SharedBuffer::destroy_if_unused()
diff --git a/Kernel/SharedBuffer.h b/Kernel/SharedBuffer.h
index 14d714ed40..acdcfe34d6 100644
--- a/Kernel/SharedBuffer.h
+++ b/Kernel/SharedBuffer.h
@@ -36,8 +36,9 @@ namespace Kernel {
 class SharedBuffer {
 private:
     struct Reference {
-        Reference(ProcessID pid)
+        Reference(ProcessID pid, unsigned count = 0)
             : pid(pid)
+            , count(count)
         {
         }
 
@@ -69,7 +70,7 @@ public:
     void share_with(ProcessID peer_pid);
     void share_globally() { m_global = true; }
     void deref_for_process(Process& process);
-    void disown(ProcessID pid);
+    bool disown(ProcessID pid);
     static void share_all_shared_buffers(Process& from_process, Process& with_process);
     size_t size() const { return m_vmobject->size(); }
     void destroy_if_unused();
diff --git a/Kernel/Syscalls/shbuf.cpp b/Kernel/Syscalls/shbuf.cpp
index 6d85c2d50a..03ff104c62 100644
--- a/Kernel/Syscalls/shbuf.cpp
+++ b/Kernel/Syscalls/shbuf.cpp
@@ -37,8 +37,12 @@ void Process::disown_all_shared_buffers()
     Vector<SharedBuffer*, 32> buffers_to_disown;
     for (auto& it : shared_buffers().resource())
         buffers_to_disown.append(it.value.ptr());
-    for (auto* shared_buffer : buffers_to_disown)
-        shared_buffer->disown(m_pid);
+    for (auto* shared_buffer : buffers_to_disown) {
+        if (shared_buffer->disown(m_pid)) {
+            shared_buffers().resource().remove(shared_buffer->id());
+            delete shared_buffer;
+        }
+    }
 }
 
 int Process::sys$shbuf_create(int size, void** buffer)