From 5e2226b161b56efdedc462fc28282c679ca84d18 Mon Sep 17 00:00:00 2001
From: Martin Michelsen <mjem@wildblue.net>
Date: Fri, 21 Jun 2019 18:36:29 -0700
Subject: [PATCH] fix send + immediate disconnect case

- this happens for e.g. authentication failures
---
 Server.cc | 51 +++++++++++++++++++++++++++++++++++++++++++++++----
 Server.hh |  6 ++++++
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/Server.cc b/Server.cc
index e0e38de6..5808fa5c 100644
--- a/Server.cc
+++ b/Server.cc
@@ -36,14 +36,28 @@ Server::WorkerThread::WorkerThread(Server* server, int worker_num) :
 }
 
 void Server::WorkerThread::disconnect_client(struct bufferevent* bev) {
-  auto client = this->bev_to_client.at(bev);
-  this->bev_to_client.erase(bev);
   {
+    auto client = this->bev_to_client.at(bev);
+    this->bev_to_client.erase(bev);
+    this->server->client_count--;
+
     rw_guard g(client->lock, true);
-    bufferevent_free(client->bev);
     client->bev = NULL;
   }
-  this->server->client_count--;
+
+  // if the output buffer is not empty, move the client into the draining pool
+  // instead of disconnecting it, to make sure all the data gets sent
+  struct evbuffer* out_buffer = bufferevent_get_output(bev);
+  if (evbuffer_get_length(out_buffer) == 0) {
+    bufferevent_free(bev);
+  } else {
+    // the callbacks will free it when all the data is sent or the client
+    // disconnects
+    bufferevent_setcb(bev, NULL,
+        Server::WorkerThread::dispatch_on_disconnecting_client_output,
+        Server::WorkerThread::dispatch_on_disconnecting_client_error, this);
+    bufferevent_disable(bev, EV_READ);
+  }
 }
 
 void Server::WorkerThread::dispatch_on_listen_accept(
@@ -71,6 +85,18 @@ void Server::WorkerThread::dispatch_on_client_error(
   wt->server->on_client_error(*wt, bev, events);
 }
 
+void Server::WorkerThread::dispatch_on_disconnecting_client_output(
+    struct bufferevent *bev, void *ctx) {
+  WorkerThread* wt = (WorkerThread*)ctx;
+  wt->server->on_disconnecting_client_output(*wt, bev);
+}
+
+void Server::WorkerThread::dispatch_on_disconnecting_client_error(
+    struct bufferevent *bev, short events, void *ctx) {
+  WorkerThread* wt = (WorkerThread*)ctx;
+  wt->server->on_disconnecting_client_error(*wt, bev, events);
+}
+
 void Server::WorkerThread::dispatch_check_for_thread_exit(
     evutil_socket_t fd, short what, void* ctx) {
   WorkerThread* wt = (WorkerThread*)ctx;
@@ -151,6 +177,11 @@ void Server::on_client_input(Server::WorkerThread& wt,
   }
 }
 
+void Server::on_disconnecting_client_output(Server::WorkerThread& wt,
+    struct bufferevent *bev) {
+  bufferevent_free(bev);
+}
+
 void Server::on_client_error(Server::WorkerThread& wt,
     struct bufferevent *bev, short events) {
   shared_ptr<Client> c;
@@ -171,6 +202,18 @@ void Server::on_client_error(Server::WorkerThread& wt,
   }
 }
 
+void Server::on_disconnecting_client_error(Server::WorkerThread& wt,
+    struct bufferevent *bev, short events) {
+  if (events & BEV_EVENT_ERROR) {
+    int err = EVUTIL_SOCKET_ERROR();
+    log(WARNING, "[Server] disconnecting client caused %d (%s)\n", err,
+        evutil_socket_error_to_string(err));
+  }
+  if (events & (BEV_EVENT_EOF | BEV_EVENT_ERROR)) {
+    bufferevent_free(bev);
+  }
+}
+
 void Server::check_for_thread_exit(Server::WorkerThread& wt,
     evutil_socket_t fd, short what) {
   if (this->should_exit) {
diff --git a/Server.hh b/Server.hh
index 6f6f4136..b3c0a019 100644
--- a/Server.hh
+++ b/Server.hh
@@ -34,6 +34,10 @@ private:
     static void dispatch_on_client_input(struct bufferevent* bev, void* ctx);
     static void dispatch_on_client_error(struct bufferevent* bev, short events,
         void* ctx);
+    static void dispatch_on_disconnecting_client_output(struct bufferevent* bev,
+        void* ctx);
+    static void dispatch_on_disconnecting_client_error(struct bufferevent* bev,
+        short events, void* ctx);
     static void dispatch_check_for_thread_exit(evutil_socket_t fd, short what, void* ctx);
   };
 
@@ -49,6 +53,8 @@ private:
   void on_listen_error(WorkerThread& wt, struct evconnlistener *listener);
   void on_client_input(WorkerThread& wt, struct bufferevent *bev);
   void on_client_error(WorkerThread& wt, struct bufferevent *bev, short events);
+  void on_disconnecting_client_output(WorkerThread& wt, struct bufferevent *bev);
+  void on_disconnecting_client_error(WorkerThread& wt, struct bufferevent *bev, short events);
   void check_for_thread_exit(WorkerThread& wt, evutil_socket_t fd, short what);
 
   void receive_and_process_commands(std::shared_ptr<Client> c, struct bufferevent* buf);