Revert "Revert "Connection and websockets fixes""

This reverts commit a8086647b103f55116ce4c872e1455ebf1f3e346.

Reason for revert: Restoring commit c00500b as base for upload image issue fix

Change-Id: I1dd5d3fda2d1ee6f4027193a0506d5ca764b01e4
Signed-off-by: Jan Sowinski <jan.sowinski@intel.com>
diff --git a/http/http_connection.h b/http/http_connection.h
index 7d24fe7..4ef3bc6 100644
--- a/http/http_connection.h
+++ b/http/http_connection.h
@@ -248,7 +248,8 @@
     1024 * 1024 * BMCWEB_HTTP_REQ_BODY_LIMIT_MB;
 
 template <typename Adaptor, typename Handler, typename... Middlewares>
-class Connection
+class Connection : public std::enable_shared_from_this<
+                       Connection<Adaptor, Handler, Middlewares...>>
 {
   public:
     Connection(boost::asio::io_context& ioService, Handler* handlerIn,
@@ -474,16 +475,15 @@
                                      boost::beast::ssl_stream<
                                          boost::asio::ip::tcp::socket>>)
         {
-            adaptor.async_handshake(
-                boost::asio::ssl::stream_base::server,
-                [this](const boost::system::error_code& ec) {
-                    if (ec)
-                    {
-                        checkDestroy();
-                        return;
-                    }
-                    doReadHeaders();
-                });
+            adaptor.async_handshake(boost::asio::ssl::stream_base::server,
+                                    [this, self(shared_from_this())](
+                                        const boost::system::error_code& ec) {
+                                        if (ec)
+                                        {
+                                            return;
+                                        }
+                                        doReadHeaders();
+                                    });
         }
         else
         {
@@ -561,18 +561,21 @@
 
             if (!res.completed)
             {
+                needToCallAfterHandlers = true;
+                res.completeRequestHandler = [self(shared_from_this())] {
+                    self->completeRequest();
+                };
                 if (req->isUpgrade() &&
                     boost::iequals(
                         req->getHeaderValue(boost::beast::http::field::upgrade),
                         "websocket"))
                 {
                     handler->handleUpgrade(*req, res, std::move(adaptor));
+                    // delete lambda with self shared_ptr
+                    // to enable connection destruction
+                    res.completeRequestHandler = nullptr;
                     return;
                 }
-                res.completeRequestHandler = [this] {
-                    this->completeRequest();
-                };
-                needToCallAfterHandlers = true;
                 handler->handle(*req, res);
             }
             else
@@ -638,15 +641,16 @@
                 *middlewares, ctx, *req, res);
         }
 
-        // auto self = this->shared_from_this();
-        res.completeRequestHandler = res.completeRequestHandler = [] {};
-
         if (!isAlive())
         {
             // BMCWEB_LOG_DEBUG << this << " delete (socket is closed) " <<
             // isReading
             // << ' ' << isWriting;
             // delete this;
+
+            // delete lambda with self shared_ptr
+            // to enable connection destruction
+            res.completeRequestHandler = nullptr;
             return;
         }
         if (res.body().empty() && !res.jsonValue.empty())
@@ -683,21 +687,23 @@
         res.keepAlive(req->keepAlive());
 
         doWrite();
+
+        // delete lambda with self shared_ptr
+        // to enable connection destruction
+        res.completeRequestHandler = nullptr;
     }
 
   private:
     void doReadHeaders()
     {
-        // auto self = this->shared_from_this();
-        isReading = true;
         BMCWEB_LOG_DEBUG << this << " doReadHeaders";
 
         // Clean up any previous Connection.
         boost::beast::http::async_read_header(
             adaptor, buffer, *parser,
-            [this](const boost::system::error_code& ec,
-                   std::size_t bytes_transferred) {
-                isReading = false;
+            [this,
+             self(shared_from_this())](const boost::system::error_code& ec,
+                                       std::size_t bytes_transferred) {
                 BMCWEB_LOG_ERROR << this << " async_read_header "
                                  << bytes_transferred << " Bytes";
                 bool errorWhileReading = false;
@@ -722,7 +728,6 @@
                     cancelDeadlineTimer();
                     close();
                     BMCWEB_LOG_DEBUG << this << " from read(1)";
-                    checkDestroy();
                     return;
                 }
 
@@ -740,17 +745,15 @@
 
     void doRead()
     {
-        // auto self = this->shared_from_this();
-        isReading = true;
         BMCWEB_LOG_DEBUG << this << " doRead";
 
         boost::beast::http::async_read(
             adaptor, buffer, *parser,
-            [this](const boost::system::error_code& ec,
-                   std::size_t bytes_transferred) {
+            [this,
+             self(shared_from_this())](const boost::system::error_code& ec,
+                                       std::size_t bytes_transferred) {
                 BMCWEB_LOG_DEBUG << this << " async_read " << bytes_transferred
                                  << " Bytes";
-                isReading = false;
 
                 bool errorWhileReading = false;
                 if (ec)
@@ -771,7 +774,6 @@
                     cancelDeadlineTimer();
                     close();
                     BMCWEB_LOG_DEBUG << this << " from read(1)";
-                    checkDestroy();
                     return;
                 }
                 handle();
@@ -780,30 +782,26 @@
 
     void doWrite()
     {
-        // auto self = this->shared_from_this();
-        isWriting = true;
         BMCWEB_LOG_DEBUG << this << " doWrite";
         res.preparePayload();
         serializer.emplace(*res.stringResponse);
         boost::beast::http::async_write(
             adaptor, *serializer,
-            [&](const boost::system::error_code& ec,
-                std::size_t bytes_transferred) {
-                isWriting = false;
+            [this,
+             self(shared_from_this())](const boost::system::error_code& ec,
+                                       std::size_t bytes_transferred) {
                 BMCWEB_LOG_DEBUG << this << " async_write " << bytes_transferred
                                  << " bytes";
 
                 if (ec)
                 {
                     BMCWEB_LOG_DEBUG << this << " from write(2)";
-                    checkDestroy();
                     return;
                 }
                 if (!res.keepAlive())
                 {
                     close();
                     BMCWEB_LOG_DEBUG << this << " from write(1)";
-                    checkDestroy();
                     return;
                 }
 
@@ -820,29 +818,25 @@
             });
     }
 
-    void checkDestroy()
-    {
-        BMCWEB_LOG_DEBUG << this << " isReading " << isReading << " isWriting "
-                         << isWriting;
-        if (!isReading && !isWriting)
-        {
-            BMCWEB_LOG_DEBUG << this << " delete (idle) ";
-            delete this;
-        }
-    }
-
     void cancelDeadlineTimer()
     {
-        BMCWEB_LOG_DEBUG << this << " timer cancelled: " << &timerQueue << ' '
-                         << timerCancelKey;
-        timerQueue.cancel(timerCancelKey);
+        if (timerCancelKey)
+        {
+            BMCWEB_LOG_DEBUG << this << " timer cancelled: " << &timerQueue
+                             << ' ' << *timerCancelKey;
+            timerQueue.cancel(*timerCancelKey);
+            timerCancelKey.reset();
+        }
     }
 
     void startDeadline()
     {
         cancelDeadlineTimer();
 
-        timerCancelKey = timerQueue.add([this] {
+        timerCancelKey = timerQueue.add([this, self(shared_from_this())] {
+            // Mark timer as not active to avoid canceling it during
+            // Connection destructor which leads to double free issue
+            timerCancelKey.reset();
             if (!isAlive())
             {
                 return;
@@ -850,7 +844,7 @@
             close();
         });
         BMCWEB_LOG_DEBUG << this << " timer added: " << &timerQueue << ' '
-                         << timerCancelKey;
+                         << *timerCancelKey;
     }
 
   private:
@@ -877,10 +871,8 @@
 
     const std::string& serverName;
 
-    size_t timerCancelKey = 0;
+    std::optional<size_t> timerCancelKey;
 
-    bool isReading{};
-    bool isWriting{};
     bool needToCallAfterHandlers{};
     bool needToStartReadAfterComplete{};
 
@@ -889,5 +881,8 @@
 
     std::function<std::string()>& getCachedDateStr;
     detail::TimerQueue& timerQueue;
+
+    using std::enable_shared_from_this<
+        Connection<Adaptor, Handler, Middlewares...>>::shared_from_this;
 };
 } // namespace crow