From 8e73b9064fdd414b9a71806a44f0e0f176b5a88d Mon Sep 17 00:00:00 2001
From: Xinnan Xie <xiexinnan@bytedance.com>
Date: Wed, 23 Aug 2023 11:14:48 +0800
Subject: kvm_websocket: Fix crash on dangling reference

Kvm_websocket captures the this pointer in the handler lambda of the
socket. When the callback is called, if the object has been destructed,
there will be a crash problem. This is fixed by using weak_from_this in
the callback, if the object was destructed, the callback just returns
without doing anything.

Tested:
1. Open two kvm sessions in WebUI, and keep refreshing in one of the
pages, there is a small chance of coredump happening.
Debug infomation shows:
```
bmcweb[5331]: DEBUG: doRead. conn:0x28d19a0. this: 0x284d470
bmcweb[5331]: DEBUG: doRead. conn:0x28d19a0. this: 0x284d470
bmcweb[5331]: DEBUG: doRead. conn:0x28d19a0. this: 0x284d470
bmcweb[5331]: DEBUG: doRead. conn:0x28d19a0. this: 0x284d470
bmcweb[5331]: DEBUG: doRead. conn:0x28d19a0. this: 0x284d470
bmcweb[5331]: DEBUG: onclose. conn:0x28d19a0
bmcweb[5331]: DEBUG: doRead. conn:0x2876648. this: 0x284d470
systemd[1]: bmeweb.service: Main process exited, code=dumped, status=11/SEGV
systemd[1]: bmcweb.service: Failed with result 'core-dump
systemd[1]: Started Start bmweb server.
```
2. After this fix no coredump occurred.

Change-Id: I7bba9b67c470def90ddb1e471a0ac95edd6165e5
Signed-off-by: Xinnan Xie <xiexinnan@bytedance.com>
---
 include/kvm_websocket.hpp | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

(limited to 'include')

diff --git a/include/kvm_websocket.hpp b/include/kvm_websocket.hpp
index e3131bb3c2..d04c19fd93 100644
--- a/include/kvm_websocket.hpp
+++ b/include/kvm_websocket.hpp
@@ -14,7 +14,7 @@ namespace obmc_kvm
 
 static constexpr const uint maxSessions = 4;
 
-class KvmSession
+class KvmSession : public std::enable_shared_from_this<KvmSession>
 {
   public:
     explicit KvmSession(crow::websocket::Connection& connIn) :
@@ -71,7 +71,13 @@ class KvmSession
                          bytes);
         hostSocket.async_read_some(
             outputBuffer.prepare(outputBuffer.capacity() - outputBuffer.size()),
-            [this](const boost::system::error_code& ec, std::size_t bytesRead) {
+            [this, weak(weak_from_this())](const boost::system::error_code& ec,
+                                           std::size_t bytesRead) {
+            auto self = weak.lock();
+            if (self == nullptr)
+            {
+                return;
+            }
             BMCWEB_LOG_DEBUG("conn:{}, read done.  Read {} bytes",
                              logPtr(&conn), bytesRead);
             if (ec)
@@ -115,9 +121,15 @@ class KvmSession
         }
 
         doingWrite = true;
-        hostSocket.async_write_some(inputBuffer.data(),
-                                    [this](const boost::system::error_code& ec,
+        hostSocket.async_write_some(
+            inputBuffer.data(),
+            [this, weak(weak_from_this())](const boost::system::error_code& ec,
                                            std::size_t bytesWritten) {
+            auto self = weak.lock();
+            if (self == nullptr)
+            {
+                return;
+            }
             BMCWEB_LOG_DEBUG("conn:{}, Wrote {}bytes", logPtr(&conn),
                              bytesWritten);
             doingWrite = false;
@@ -140,7 +152,7 @@ class KvmSession
             }
 
             doWrite();
-        });
+            });
     }
 
     crow::websocket::Connection& conn;
@@ -151,7 +163,7 @@ class KvmSession
 };
 
 using SessionMap = boost::container::flat_map<crow::websocket::Connection*,
-                                              std::unique_ptr<KvmSession>>;
+                                              std::shared_ptr<KvmSession>>;
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 static SessionMap sessions;
 
@@ -171,7 +183,7 @@ inline void requestRoutes(App& app)
                 return;
             }
 
-            sessions[&conn] = std::make_unique<KvmSession>(conn);
+            sessions[&conn] = std::make_shared<KvmSession>(conn);
         })
         .onclose([](crow::websocket::Connection& conn, const std::string&) {
             sessions.erase(&conn);
-- 
cgit v1.2.3