From 855c04eb3b452d8c1280405cd5e8c8b7d8359c71 Mon Sep 17 00:00:00 2001
From: Bruce Toll <btoll@dhsus.com>
Date: Sun, 10 Mar 2019 15:02:19 -0400
Subject: [PATCH] Fix webrequest leak with websockets due to race

There is a race at websession destruction time when flush() is called
on a websocket that may still be actively writing.

In particular, this can occur when a Wt application sends output as
part of a final quit(), e.g. to display a timeout message in a dialog.

The fix is to keep the websession alive until any outstanding
websocket write is complete. This needs to be done from the time
the websocket is created.
---
 src/http/HTTPRequest.C |  5 +++++
 src/http/HTTPRequest.h |  1 +
 src/http/WtReply.C     | 15 +++++++++++++++
 src/http/WtReply.h     |  4 ++++
 src/web/WebRequest.C   |  5 +++++
 src/web/WebRequest.h   |  3 +++
 src/web/WebSession.C   |  4 ++++
 7 files changed, 37 insertions(+)

diff --git a/src/http/HTTPRequest.C b/src/http/HTTPRequest.C
index ba53f423..524a8477 100644
--- a/src/http/HTTPRequest.C
+++ b/src/http/HTTPRequest.C
@@ -255,5 +255,10 @@ const std::vector<std::pair<std::string, std::string> > &HTTPRequest::urlParams(
   return p->request().url_params;
 }
 
+void HTTPRequest::setWebSession(WebSessionPtr webSession)
+{
+  reply_->setWebSession(webSession);
+}
+
 } // namespace server
 } // namespace http
diff --git a/src/http/HTTPRequest.h b/src/http/HTTPRequest.h
index 11b18d76..ce201c7a 100644
--- a/src/http/HTTPRequest.h
+++ b/src/http/HTTPRequest.h
@@ -56,6 +56,7 @@ public:
   bool isSynchronous() const;
   virtual Wt::WSslInfo *sslInfo() const override;
   virtual const std::vector<std::pair<std::string, std::string> > &urlParams() const override;
+  virtual void setWebSession(WebSessionPtr webSession) override;
 
 private:
   WtReplyPtr reply_;
diff --git a/src/http/WtReply.C b/src/http/WtReply.C
index 043ee84b..8b2f6d1d 100644
--- a/src/http/WtReply.C
+++ b/src/http/WtReply.C
@@ -8,6 +8,7 @@
 #include "WtReply.h"
 #include "StockReply.h"
 #include "HTTPRequest.h"
+#include "WebSession.h"
 #include "WebController.h"
 #include "Server.h"
 #include "WebUtils.h"
@@ -436,6 +437,11 @@ void WtReply::writeDone(bool success)
   out_buf_.consume(sending_);
   sending_ = 0;
 
+  if (webSessionLock_) {
+    LOG_DEBUG("WtReply::writeDone webSessionLock_ about to set null, currently: " << webSessionLock_.get());
+    webSessionLock_ = nullptr;
+  }
+
   if (fetchMoreDataCallback_) {
     Wt::WebRequest::WriteCallback f = fetchMoreDataCallback_;
     fetchMoreDataCallback_ = nullptr;
@@ -448,6 +454,9 @@ void WtReply::send(const Wt::WebRequest::WriteCallback& callBack,
 {
   LOG_DEBUG("WtReply::send(): " << sending_);
 
+  webSessionLock_ = webSessionWeakPtr_.lock();
+  LOG_DEBUG("WtReply::send set webSessionLock_ to: " << webSessionLock_.get());
+
   fetchMoreDataCallback_ = callBack;
 
   if (sending_ != 0) {
@@ -662,6 +671,12 @@ bool WtReply::nextContentBuffers(std::vector<asio::const_buffer>& result)
 }
 
 
+void WtReply::setWebSession(Wt::WebRequest::WebSessionPtr webSession)
+{
+    webSessionWeakPtr_ = webSession;
+}
+
+
 #ifdef WTHTTP_WITH_ZLIB
 
 bool WtReply::initDeflate() 
diff --git a/src/http/WtReply.h b/src/http/WtReply.h
index b55aec38..f7f68ba7 100644
--- a/src/http/WtReply.h
+++ b/src/http/WtReply.h
@@ -15,6 +15,7 @@
 #include "Reply.h"
 #include "../web/Configuration.h"
 #include "../web/WebRequest.h"
+#include "../web/WebSession.h"
 
 #ifdef WTHTTP_WITH_ZLIB
 #include <zlib.h>
@@ -59,6 +60,7 @@ public:
 	    bool responseComplete);
   void readWebSocketMessage(const Wt::WebRequest::ReadCallback& callBack);
   bool readAvailable();
+  void setWebSession(Wt::WebRequest::WebSessionPtr webSession);
 
   std::istream& in() { return *in_; }
   std::ostream& out() { return out_; }
@@ -106,6 +108,8 @@ private:
   bool initDeflate();
 
   z_stream zOutState_;
+  std::weak_ptr<Wt::WebSession> webSessionWeakPtr_;
+  Wt::WebRequest::WebSessionPtr webSessionLock_;
 #endif
 };
 
diff --git a/src/web/WebRequest.C b/src/web/WebRequest.C
index 1bb2a356..35aec6f2 100644
--- a/src/web/WebRequest.C
+++ b/src/web/WebRequest.C
@@ -337,4 +337,9 @@ const std::vector<std::pair<std::string, std::string> > &WebRequest::urlParams()
   return urlParams_;
 }
 
+void WebRequest::setWebSession(WebSessionPtr webSession)
+{
+  throw WException("should not get here");
+}
+
 }
diff --git a/src/web/WebRequest.h b/src/web/WebRequest.h
index bdc426f9..f6445b03 100644
--- a/src/web/WebRequest.h
+++ b/src/web/WebRequest.h
@@ -48,6 +48,8 @@ public:
   typedef std::function<void(WebReadEvent)> ReadCallback;
   typedef std::function<void(void)> DisconnectCallback;
 
+  typedef std::shared_ptr<Wt::WebSession> WebSessionPtr;
+
   /*
    * Signal that the response should be flushed.
    *
@@ -197,6 +199,7 @@ public:
   virtual WSslInfo *sslInfo() const = 0;
 
   virtual const std::vector<std::pair<std::string, std::string> >& urlParams() const;
+  virtual void setWebSession(WebSessionPtr webSession);
 
 protected:
   const EntryPoint *entryPoint_;
diff --git a/src/web/WebSession.C b/src/web/WebSession.C
index 5b8c756c..6318f48e 100644
--- a/src/web/WebSession.C
+++ b/src/web/WebSession.C
@@ -17,6 +17,7 @@
 #include "Wt/WServer.h"
 #include "Wt/WTimerWidget.h"
 #include "Wt/Http/Request.h"
+#include "http/HTTPRequest.h"
 
 #include "CgiParser.h"
 #include "Configuration.h"
@@ -252,8 +253,10 @@ WebSession::~WebSession()
   }
 
   if (webSocket_) {
+    LOG_DEBUG("~WebSession: before webSocket_ flush/null");
     webSocket_->flush();
     webSocket_ = nullptr;
+    LOG_DEBUG("~WebSession: after webSocket_ flush/null");
   }
 
   if (deferredResponse_) {
@@ -1776,6 +1779,7 @@ void WebSession::handleWebSocketRequest(Handler& handler)
   }
 
   webSocket_ = handler.response();
+  webSocket_->setWebSession(shared_from_this());
   canWriteWebSocket_ = false;
   webSocketConnected_ = false;
 
-- 
2.18.1

