Reworked "upgraded" closure logic: resources deallocated and sockets are closed

asynchronously only in daemon's thread and only when all data was forwarded and
application signaled about upgraded closure.

8 files changed, 324 insertions, 231 deletions

diff --git a/src/include/microhttpd.h b/src/include/microhttpd.h
index efba0ee6..a099d47f 100644
--- a/src/include/microhttpd.h
+++ b/src/include/microhttpd.h
@@ -126,7 +126,7 @@ typedef intptr_t ssize_t;
  * Current version of the library.
  * 0x01093001 = 1.9.30-1.
  */
-#define MHD_VERSION 0x00095103
+#define MHD_VERSION 0x00095204
 
 /**
  * MHD-internal return code for "YES".
diff --git a/src/microhttpd/connection.c b/src/microhttpd/connection.c
index 52b67b05..d3ccdf26 100644
--- a/src/microhttpd/connection.c
+++ b/src/microhttpd/connection.c
@@ -22,6 +22,7 @@
  * @brief  Methods for managing connections
  * @author Daniel Pittman
  * @author Christian Grothoff
+ * @author Karlson2k (Evgeny Grin)
  */
 
 #include "internal.h"
@@ -527,6 +528,63 @@ MHD_connection_close_ (struct MHD_Connection *connection,
 
 
 /**
+ * Stop TLS forwarding on upgraded connection and
+ * reflect remote disconnect state to socketpair.
+ * @remark In thread-per-connection mode this function
+ * can be called from any thread, in other modes this
+ * function must be called only from thread that process
+ * daemon's select()/poll()/etc.
+ *
+ * @param connection the upgraded connection
+ */
+void
+MHD_connection_finish_forward_ (struct MHD_Connection *connection)
+{
+  struct MHD_Daemon *daemon = connection->daemon;
+  struct MHD_UpgradeResponseHandle *urh = connection->urh;
+
+  if (0 == (daemon->options & MHD_USE_TLS))
+    return; /* Nothing to do with non-TLS connection. */
+
+  if (0 == (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
+    DLL_remove (daemon->urh_head,
+                daemon->urh_tail,
+                urh);
+#if EPOLL_SUPPORT
+  if ( (0 != (daemon->options & MHD_USE_EPOLL)) &&
+       (0 != epoll_ctl (daemon->epoll_upgrade_fd,
+                        EPOLL_CTL_DEL,
+                        connection->socket_fd,
+                        NULL)) )
+    {
+      MHD_PANIC (_("Failed to remove FD from epoll set\n"));
+    }
+#endif /* EPOLL_SUPPORT */
+  if (MHD_INVALID_SOCKET != urh->mhd.socket)
+    {
+#if EPOLL_SUPPORT
+      if ( (0 != (daemon->options & MHD_USE_EPOLL)) &&
+           (0 != epoll_ctl (daemon->epoll_upgrade_fd,
+                            EPOLL_CTL_DEL,
+                            urh->mhd.socket,
+                            NULL)) )
+        {
+          MHD_PANIC (_("Failed to remove FD from epoll set\n"));
+        }
+#endif /* EPOLL_SUPPORT */
+      /* Reflect remote disconnect to application by breaking
+       * socketpair connection. */
+      shutdown (urh->mhd.socket, SHUT_RDWR);
+    }
+  /* Socketpair sockets will remain open as they will be
+   * used with MHD_UPGRADE_ACTION_CLOSE. They will be
+   * closed by MHD_cleanup_upgraded_connection_() during
+   * connection's final cleanup.
+   */
+}
+
+
+/**
  * A serious error occured, close the
  * connection (and notify the application).
  *
@@ -3080,8 +3138,9 @@ MHD_connection_handle_idle (struct MHD_Connection *connection)
           cleanup_connection (connection);
           return MHD_NO;
         case MHD_CONNECTION_UPGRADE:
-        case MHD_CONNECTION_UPGRADE_CLOSED:
           return MHD_YES; /* keep open */
+        case MHD_CONNECTION_UPGRADE_CLOSED:
+          return MHD_YES; /* "Upgraded" connection should be closed in special way. */
         default:
           EXTRA_CHECK (0);
           break;
diff --git a/src/microhttpd/connection.h b/src/microhttpd/connection.h
index 0fce7346..a962a40c 100644
--- a/src/microhttpd/connection.h
+++ b/src/microhttpd/connection.h
@@ -22,6 +22,7 @@
  * @brief  Methods for managing connections
  * @author Daniel Pittman
  * @author Christian Grothoff
+ * @author Karlson2k (Evgeny Grin)
  */
 
 #ifndef CONNECTION_H
@@ -97,6 +98,15 @@ MHD_connection_close_ (struct MHD_Connection *connection,
                        enum MHD_RequestTerminationCode termination_code);
 
 
+/**
+ * Stop TLS forwarding on upgraded connection and
+ * reflect remote disconnect state to socketpair.
+ * @param connection the upgraded connection
+ */
+void
+MHD_connection_finish_forward_ (struct MHD_Connection *connection);
+
+
 #ifdef EPOLL_SUPPORT
 /**
  * Perform epoll processing, possibly moving the connection back into
diff --git a/src/microhttpd/daemon.c b/src/microhttpd/daemon.c
index 2bce6d15..622ab34e 100644
--- a/src/microhttpd/daemon.c
+++ b/src/microhttpd/daemon.c
@@ -23,6 +23,7 @@
  * @brief  A minimal-HTTP server library
  * @author Daniel Pittman
  * @author Christian Grothoff
+ * @author Karlson2k (Evgeny Grin)
  */
 #include "platform.h"
 #include "mhd_threads.h"
@@ -688,6 +689,7 @@ urh_to_fdset (struct MHD_UpgradeResponseHandle *urh,
                               fd_setsize)) )
     return MHD_NO;
   if ( (0 == (MHD_EPOLL_STATE_WRITE_READY & urh->mhd.celi)) &&
+       (MHD_NO == urh->was_closed) &&
        (MHD_INVALID_SOCKET != urh->mhd.socket) &&
        (! MHD_add_to_fd_set_ (urh->mhd.socket,
                               ws,
@@ -695,6 +697,7 @@ urh_to_fdset (struct MHD_UpgradeResponseHandle *urh,
                               fd_setsize)) )
     return MHD_NO;
   if ( (urh->in_buffer_used < urh->in_buffer_size) &&
+       (MHD_NO == urh->was_closed) &&
        (MHD_INVALID_SOCKET != urh->connection->socket_fd) &&
        (! MHD_add_to_fd_set_ (urh->connection->socket_fd,
                               rs,
@@ -934,73 +937,19 @@ call_handlers (struct MHD_Connection *con,
 void
 MHD_cleanup_upgraded_connection_ (struct MHD_Connection *connection)
 {
-  struct MHD_Daemon *daemon = connection->daemon;
   struct MHD_UpgradeResponseHandle *urh = connection->urh;
 
-#if HTTPS_SUPPORT
-  if ( (NULL != urh) &&
-       (0 != (daemon->options & MHD_USE_TLS)) )
-    {
-      if (0 == (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
-        DLL_remove (daemon->urh_head,
-                    daemon->urh_tail,
-                    urh);
-#if EPOLL_SUPPORT
-      if (0 != (daemon->options & MHD_USE_EPOLL))
-        {
-          /* epoll documentation suggests that closing a FD
-             automatically removes it from the epoll set; however,
-             this is not true as if we fail to do manually remove it,
-             we are still seeing an event for this fd in epoll,
-             causing grief (use-after-free...) --- at least on my
-             system. */
-          if (0 != epoll_ctl (daemon->epoll_upgrade_fd,
-                              EPOLL_CTL_DEL,
-                              connection->socket_fd,
-                              NULL))
-            MHD_PANIC (_("Failed to remove FD from epoll set\n"));
-        }
-#endif /* EPOLL_SUPPORT */
-      if (MHD_INVALID_SOCKET != urh->mhd.socket)
-        {
-          /* epoll documentation suggests that closing a FD
-             automatically removes it from the epoll set; however,
-             this is not true as if we fail to do manually remove it,
-             we are still seeing an event for this fd in epoll,
-             causing grief (use-after-free...) --- at least on my
-             system. */
-#if EPOLL_SUPPORT
-          if ( (0 != (daemon->options & MHD_USE_EPOLL)) &&
-               (0 != epoll_ctl (daemon->epoll_upgrade_fd,
-                                EPOLL_CTL_DEL,
-                                urh->mhd.socket,
-                                NULL)) )
-            MHD_PANIC (_("Failed to remove FD from epoll set\n"));
-#endif /* EPOLL_SUPPORT */
-          MHD_socket_close_chk_ (urh->mhd.socket);
-        }
-      if (MHD_INVALID_SOCKET != urh->app.socket)
-        MHD_socket_close_chk_ (urh->app.socket);
-    }
-#endif /* HTTPS_SUPPORT */
-  if (0 == (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
-    {
-      /* resuming the connection will indirectly close it */
-      MHD_resume_connection (connection);
-#if HTTPS_SUPPORT
-      if (0 != (daemon->options & MHD_USE_TLS))
-        {
-          gnutls_bye (connection->tls_session,
-                      GNUTLS_SHUT_RDWR);
-        }
-#endif
-    }
-  else
-    {
-      /* the thread would no longer close it */
-      MHD_connection_close_ (connection,
-                             MHD_REQUEST_TERMINATED_COMPLETED_OK);
-    }
+  /* Signal remote client the end of TLS connection by
+   * gracefully closing TLS session. */
+  if (0 != (connection->daemon->options & MHD_USE_TLS))
+    gnutls_bye (connection->tls_session,
+                GNUTLS_SHUT_WR);
+
+  if (MHD_INVALID_SOCKET != urh->mhd.socket)
+    MHD_socket_close_chk_ (urh->mhd.socket);
+
+  if (MHD_INVALID_SOCKET != urh->app.socket)
+    MHD_socket_close_chk_ (urh->app.socket);
 
   connection->urh = NULL;
   if (NULL != urh)
@@ -1259,14 +1208,6 @@ process_urh (struct MHD_UpgradeResponseHandle *urh)
       urh->app.celi &= ~MHD_EPOLL_STATE_WRITE_READY;
       urh->mhd.celi &= ~MHD_EPOLL_STATE_READ_READY;
     }
-
-  /* cleanup connection if it was closed and all data was sent */
-  if ( (MHD_NO != urh->was_closed) &&
-       (0 == urh->out_buffer_size) &&
-       (0 == urh->out_buffer_used) )
-    {
-      MHD_cleanup_upgraded_connection_ (urh->connection);
-    }
 }
 #endif
 
@@ -1292,7 +1233,8 @@ thread_main_connection_upgrade (struct MHD_Connection *con)
   if ( (0 != (daemon->options & MHD_USE_TLS)) &&
       (0 == (daemon->options & MHD_USE_POLL)))
     {
-      while (MHD_CONNECTION_UPGRADE == con->state)
+      while ( (MHD_CONNECTION_UPGRADE == con->state) ||
+              (0 != urh->out_buffer_used) )
         {
           /* use select */
           fd_set rs;
@@ -1356,7 +1298,8 @@ thread_main_connection_upgrade (struct MHD_Connection *con)
       /* use poll() */
       const unsigned int timeout = UINT_MAX;
 
-      while (MHD_CONNECTION_UPGRADE == con->state)
+      while ( (MHD_CONNECTION_UPGRADE == con->state) ||
+              (0 != urh->out_buffer_used) )
         {
           struct pollfd p[2];
 
@@ -1410,6 +1353,10 @@ thread_main_connection_upgrade (struct MHD_Connection *con)
 #endif
   /* end HTTPS */
 #endif
+  /* TLS forwarding was finished. Cleanup socketpair. */
+  MHD_connection_finish_forward_ (con);
+  /* Do not set 'urh->clean_ready' yet as 'urh' will be used
+   * in connection thread for a little while. */
 }
 
 
@@ -1456,7 +1403,7 @@ thread_main_handle_connection (void *data)
       const unsigned int timeout = daemon->connection_timeout;
       _MHD_bool was_suspended = 0;
 
-      if (MHD_NO != con->suspended)
+      if (MHD_NO != con->suspended && NULL == con->urh)
         {
           /* Connection was suspended, wait for resume. */
           was_suspended = !0;
@@ -1721,7 +1668,11 @@ thread_main_handle_connection (void *data)
             goto exit;
         }
 #endif
-      if (MHD_CONNECTION_UPGRADE == con->state)
+      /* Check for 'MHD_CONNECTION_UPGRADE_CLOSED' too:
+       * application can finish with "upgraded" connection
+       * before this thread process it for the first time. */
+      if ( (MHD_CONNECTION_UPGRADE == con->state) ||
+           (MHD_CONNECTION_UPGRADE_CLOSED == con->state) )
         {
           /* Normal HTTP processing is finished,
            * notify application. */
@@ -1734,13 +1685,17 @@ thread_main_handle_connection (void *data)
           con->client_aware = MHD_NO;
 
           thread_main_connection_upgrade (con);
-#if HTTPS_SUPPORT
-          if (0 != (daemon->options & MHD_USE_TLS) )
-            break;
-#endif
+          /* MHD_connection_finish_forward_() was called by thread_main_connection_upgrade(). */
+
+          /* "Upgraded" data will not be used in this thread from this point. */
+          con->urh->clean_ready = MHD_YES;
+          /* If 'urh->was_closed' set to MHD_YES, connection will be
+           * moved immediately to cleanup list. Otherwise connection
+           * will stay in suspended list until 'urh' will be marked
+           * with 'was_closed' by application. */
+          MHD_resume_connection(con);
 
-          /* skip usual clean up EXCEPT for the completion
-             notification (which must be done in this thread)! */
+          /* skip usual clean up  */
           return (MHD_THRD_RTRN_TYPE_) 0;
         }
     }
@@ -2459,6 +2414,8 @@ MHD_resume_connection (struct MHD_Connection *connection)
 /**
  * Run through the suspended connections and move any that are no
  * longer suspended back to the active state.
+ * @remark To be called only from thread that process
+ * daemon's select()/poll()/etc.
  *
  * @param daemon daemon context
  * @return #MHD_YES if a connection was actually resumed
@@ -2490,40 +2447,55 @@ resume_suspended_connections (struct MHD_Daemon *daemon)
   while (NULL != (pos = next))
     {
       next = pos->next;
-      if (MHD_NO == pos->resuming)
+      if ( (MHD_NO == pos->resuming) ||
+           ((NULL != pos->urh) &&
+            ((MHD_NO == pos->urh->was_closed) || (MHD_NO == pos->urh->clean_ready))) )
         continue;
       ret = MHD_YES;
       DLL_remove (daemon->suspended_connections_head,
                   daemon->suspended_connections_tail,
                   pos);
-      DLL_insert (daemon->connections_head,
-                  daemon->connections_tail,
-                  pos);
-      if (0 == (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
+      if (NULL == pos->urh)
         {
-          if (pos->connection_timeout == daemon->connection_timeout)
-            XDLL_insert (daemon->normal_timeout_head,
-                         daemon->normal_timeout_tail,
-                         pos);
-          else
-            XDLL_insert (daemon->manual_timeout_head,
-                         daemon->manual_timeout_tail,
-                         pos);
-        }
+          DLL_insert (daemon->connections_head,
+                      daemon->connections_tail,
+                      pos);
+          if (0 == (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
+            {
+              if (pos->connection_timeout == daemon->connection_timeout)
+                XDLL_insert (daemon->normal_timeout_head,
+                             daemon->normal_timeout_tail,
+                             pos);
+              else
+                XDLL_insert (daemon->manual_timeout_head,
+                             daemon->manual_timeout_tail,
+                             pos);
+            }
 #ifdef EPOLL_SUPPORT
-      if (0 != (daemon->options & MHD_USE_EPOLL))
+          if (0 != (daemon->options & MHD_USE_EPOLL))
+            {
+              if (0 != (pos->epoll_state & MHD_EPOLL_STATE_IN_EREADY_EDLL))
+                MHD_PANIC ("Resumed connection was already in EREADY set\n");
+              /* we always mark resumed connections as ready, as we
+                 might have missed the edge poll event during suspension */
+              EDLL_insert (daemon->eready_head,
+                           daemon->eready_tail,
+                           pos);
+              pos->epoll_state |= MHD_EPOLL_STATE_IN_EREADY_EDLL;
+              pos->epoll_state &= ~MHD_EPOLL_STATE_SUSPENDED;
+            }
+#endif
+        }
+      else
         {
-          if (0 != (pos->epoll_state & MHD_EPOLL_STATE_IN_EREADY_EDLL))
-            MHD_PANIC ("Resumed connection was already in EREADY set\n");
-          /* we always mark resumed connections as ready, as we
-             might have missed the edge poll event during suspension */
-          EDLL_insert (daemon->eready_head,
-                       daemon->eready_tail,
-                       pos);
-          pos->epoll_state |= MHD_EPOLL_STATE_IN_EREADY_EDLL;
-          pos->epoll_state &= ~MHD_EPOLL_STATE_SUSPENDED;
+          /* Data forwarding was finished (for TLS connections) AND
+           * application was closed upgraded connection.
+           * Insert connection into cleanup list. */
+          DLL_insert (daemon->cleanup_head,
+                      daemon->cleanup_tail,
+                      pos);
+
         }
-#endif
       pos->suspended = MHD_NO;
       pos->resuming = MHD_NO;
     }
@@ -2735,6 +2707,8 @@ MHD_accept_connection (struct MHD_Daemon *daemon)
  * Free resources associated with all closed connections.
  * (destroy responses, free buffers, etc.).  All closed
  * connections are kept in the "cleanup" doubly-linked list.
+ * @remark To be called only from thread that
+ * process daemon's select()/poll()/etc.
  *
  * @param daemon daemon to clean up
  */
@@ -2758,6 +2732,8 @@ MHD_cleanup_connections (struct MHD_Daemon *daemon)
            (MHD_NO == pos->thread_joined) &&
            (! MHD_join_thread_ (pos->pid)) )
         MHD_PANIC (_("Failed to join a thread\n"));
+      if (NULL != pos->urh)
+        MHD_cleanup_upgraded_connection_ (pos);
       MHD_pool_destroy (pos->pool);
 #if HTTPS_SUPPORT
       if (NULL != pos->tls_session)
@@ -3014,6 +2990,17 @@ MHD_run_from_select (struct MHD_Daemon *daemon,
                       write_fd_set);
       /* call generic forwarding function for passing data */
       process_urh (urh);
+      /* Finished forwarding? */
+      if ( (0 == urh->in_buffer_size) &&
+           (0 == urh->out_buffer_size) &&
+           (0 == urh->in_buffer_used) &&
+           (0 == urh->out_buffer_used) )
+        {
+          MHD_connection_finish_forward_ (urh->connection);
+          urh->clean_ready = MHD_YES;
+          /* Resuming will move connection to cleanup list. */
+          MHD_resume_connection(urh->connection);
+        }
     }
 #endif
   MHD_cleanup_connections (daemon);
@@ -3367,28 +3354,45 @@ MHD_poll_all (struct MHD_Daemon *daemon,
       {
         if (i >= num_connections)
           break; /* connection list changed somehow, retry later ... */
+
+        /* Get next connection here as connection can be removed
+         * from 'daemon->urh_head' list. */
         urhn = urh->next;
-        if (p[poll_server+i].fd != urh->connection->socket_fd)
-          continue; /* fd mismatch, something else happened, retry later ... */
-        if (0 != (p[poll_server+i].revents & POLLIN))
-          urh->app.celi |= MHD_EPOLL_STATE_READ_READY;
-        if (0 != (p[poll_server+i].revents & POLLOUT))
-          urh->app.celi |= MHD_EPOLL_STATE_WRITE_READY;
+        /* Check for fd mismatch. FIXME: required for safety? */
+        if (p[poll_server+i].fd == urh->connection->socket_fd)
+          {
+            if (0 != (p[poll_server+i].revents & POLLIN))
+              urh->app.celi |= MHD_EPOLL_STATE_READ_READY;
+            if (0 != (p[poll_server+i].revents & POLLOUT))
+              urh->app.celi |= MHD_EPOLL_STATE_WRITE_READY;
+          }
         i++;
-        if (p[poll_server+i].fd != urh->mhd.socket)
+        /* Check for fd mismatch. FIXME: required for safety? */
+        if (p[poll_server+i].fd == urh->mhd.socket)
           {
-            /* fd mismatch, something else happened, retry later ... */
-            /* may still be able to do something based on updates
-               to socket_fd availability */
-            process_urh (urh);
-            continue;
+            if (0 != (p[poll_server+i].revents & POLLIN))
+              urh->mhd.celi |= MHD_EPOLL_STATE_READ_READY;
+            if (0 != (p[poll_server+i].revents & POLLOUT))
+              urh->mhd.celi |= MHD_EPOLL_STATE_WRITE_READY;
           }
-        if (0 != (p[poll_server+i].revents & POLLIN))
-          urh->mhd.celi |= MHD_EPOLL_STATE_READ_READY;
-        if (0 != (p[poll_server+i].revents & POLLOUT))
-          urh->mhd.celi |= MHD_EPOLL_STATE_WRITE_READY;
         i++;
         process_urh (urh);
+        /* Finished forwarding? */
+        if ( (0 == urh->in_buffer_size) &&
+             (0 == urh->out_buffer_size) &&
+             (0 == urh->in_buffer_used) &&
+             (0 == urh->out_buffer_used) )
+          {
+            /* MHD_connection_finish_forward_() will remove connection from
+             * 'daemon->urh_head' list. */
+            MHD_connection_finish_forward_ (urh->connection);
+            urh->clean_ready = MHD_YES;
+            /* If 'urh->was_closed' set to MHD_YES, connection will be
+             * moved immediately to cleanup list. Otherwise connection
+             * will stay in suspended list until 'urh' will be marked
+             * with 'was_closed' by application. */
+            MHD_resume_connection(urh->connection);
+          }
       }
 #endif
     /* handle 'listen' FD */
@@ -3530,12 +3534,11 @@ run_epoll_for_upgrade (struct MHD_Daemon *daemon)
 {
   struct epoll_event events[MAX_EVENTS];
   int num_events;
-  unsigned int i;
-  unsigned int j;
 
   num_events = MAX_EVENTS;
   while (MAX_EVENTS == num_events)
     {
+      unsigned int i;
       /* update event masks */
       num_events = epoll_wait (daemon->epoll_upgrade_fd,
                                events,
@@ -3555,35 +3558,8 @@ run_epoll_for_upgrade (struct MHD_Daemon *daemon)
         }
       for (i=0;i<(unsigned int) num_events;i++)
         {
-          struct UpgradeEpollHandle *ueh = events[i].data.ptr;
-          struct MHD_UpgradeResponseHandle *urh;
-
-          if (NULL == ueh)
-            continue; /* was killed, see below */
-          urh = ueh->urh;
-
-          /* In case we get two events for the same upgrade handle,
-             squash them together (otherwise the first one may
-             cause us to free the 'urh', and the second one then
-             causes a use-after-free).  */
-          for (j=i+1;j< (unsigned int) num_events;j++)
-            {
-              struct UpgradeEpollHandle *uehj = events[j].data.ptr;
-              struct MHD_UpgradeResponseHandle *urhj;
-
-              if (NULL == uehj)
-                continue; /* was killed, see below */
-              urhj = uehj->urh;
-
-              if (urh == urhj) /* yep, indeed the same! */
-                {
-                  if (0 != (events[j].events & EPOLLIN))
-                    uehj->celi |= MHD_EPOLL_STATE_READ_READY;
-                  if (0 != (events[j].events & EPOLLOUT))
-                    uehj->celi |= MHD_EPOLL_STATE_WRITE_READY;
-                }
-              events[j].data.ptr = NULL; /* kill this one */
-            }
+          struct UpgradeEpollHandle * const ueh = events[i].data.ptr;
+          struct MHD_UpgradeResponseHandle * const urh = ueh->urh;
 
           /* Update our state based on what is ready according to epoll() */
           if (0 != (events[i].events & EPOLLIN))
@@ -3591,8 +3567,21 @@ run_epoll_for_upgrade (struct MHD_Daemon *daemon)
           if (0 != (events[i].events & EPOLLOUT))
             ueh->celi |= MHD_EPOLL_STATE_WRITE_READY;
 
-          /* shuffle data based on buffers and FD readyness */
           process_urh (urh);
+          /* Finished forwarding? */
+          if ( (0 == urh->in_buffer_size) &&
+               (0 == urh->out_buffer_size) &&
+               (0 == urh->in_buffer_used) &&
+               (0 == urh->out_buffer_used) )
+            {
+              MHD_connection_finish_forward_ (urh->connection);
+              urh->clean_ready = MHD_YES;
+              /* If 'urh->was_closed' set to MHD_YES, connection will be
+               * moved immediately to cleanup list. Otherwise connection
+               * will stay in suspended list until 'urh' will be marked
+               * with 'was_closed' by application. */
+              MHD_resume_connection(urh->connection);
+            }
         }
     }
   return MHD_YES;
@@ -5403,18 +5392,23 @@ close_all_connections (struct MHD_Daemon *daemon)
 #ifdef HTTPS_SUPPORT
   struct MHD_UpgradeResponseHandle *urh;
   struct MHD_UpgradeResponseHandle *urhn;
+  const _MHD_bool used_tls = (0 != (daemon->options & MHD_USE_TLS));
+  const _MHD_bool used_thr_p_c = (0 != (daemon->options & MHD_USE_THREAD_PER_CONNECTION));
 #endif /* HTTPS_SUPPORT */
 
+#ifdef HTTPS_SUPPORT
   /* give upgraded HTTPS connections a chance to finish */
-#if HTTPS_SUPPORT
+  /* 'daemon->urh_head' is not used in thread-per-connection mode. */
   for (urh = daemon->urh_head; NULL != urh; urh = urhn)
     {
       urhn = urh->next;
       /* call generic forwarding function for passing data
-         with chance to detect that application is done;
-         fake read readyness just to be sure. */
-      urh->mhd.celi |= MHD_EPOLL_STATE_READ_READY;
+         with chance to detect that application is done. */
       process_urh (urh);
+      MHD_connection_finish_forward_ (urh->connection);
+      urh->clean_ready = MHD_YES;
+      /* Resuming will move connection to cleanup list. */
+      MHD_resume_connection(urh->connection);
     }
 #endif
 
@@ -5431,6 +5425,23 @@ close_all_connections (struct MHD_Daemon *daemon)
      traverse DLLs in peace... */
   if (0 != (daemon->options & MHD_USE_THREAD_PER_CONNECTION))
     MHD_mutex_lock_chk_ (&daemon->cleanup_connection_mutex);
+#ifdef HTTPS_SUPPORT
+  if (used_tls && used_thr_p_c)
+    {
+      struct MHD_Connection * susp;
+
+      susp = daemon->suspended_connections_head;
+      while (NULL != susp)
+        {
+          if (NULL == susp->urh) /* "Upgraded" connection? */
+            MHD_PANIC (_("MHD_stop_daemon() called while we have suspended connections.\n"));
+          else if (MHD_NO == susp->urh->clean_ready)
+            shutdown (urh->app.socket, SHUT_RDWR); /* Wake thread by shutdown of app socket. */
+          susp = susp->next;
+        }
+    }
+  else /* This 'else' is combined with next 'if'. */
+#endif /* HTTPS_SUPPORT */
   if (NULL != daemon->suspended_connections_head)
     MHD_PANIC (_("MHD_stop_daemon() called while we have suspended connections.\n"));
   for (pos = daemon->connections_head; NULL != pos; pos = pos->next)
@@ -5465,6 +5476,17 @@ close_all_connections (struct MHD_Daemon *daemon)
         pos = pos->next;
       }
     }
+
+#ifdef HTTPS_SUPPORT
+  /* Finished threads with "upgraded" connections need to be moved
+   * to cleanup list by resume_suspended_connections(). */
+  if (used_tls && used_thr_p_c)
+    {
+      daemon->resuming = MHD_YES; /* Force check for pending resume. */
+      resume_suspended_connections (daemon);
+    }
+#endif /* HTTPS_SUPPORT */
+
   /* now that we're alone, move everyone to cleanup */
   while (NULL != (pos = daemon->connections_head))
   {
diff --git a/src/microhttpd/internal.h b/src/microhttpd/internal.h
index 1f239f10..bb9974a9 100644
--- a/src/microhttpd/internal.h
+++ b/src/microhttpd/internal.h
@@ -1056,8 +1056,37 @@ struct MHD_UpgradeResponseHandle
   /**
    * Set to #MHD_YES after the application finished with the socket
    * by #MHD_UPGRADE_ACTION_CLOSE.
+   *
+   * When BOTH @e was_closed (changed by command from application)
+   * AND @e clean_ready (changed internally by MHD) are set to
+   * #MHD_YES, function #MHD_resume_connection() will move this
+   * connection to cleanup list.
+   * @remark This flag could be changed from any thread.
    */
   int was_closed;
+
+  /**
+   * Set to #MHD_YES if connection is ready for cleanup.
+   *
+   * In TLS mode functions #MHD_connection_finish_forward_() must
+   * be called before setting this flag to #MHD_YES.
+   *
+   * In thread-per-connection mode #MHD_YES in this flag means
+   * that connection's thread exited or about to exit and will
+   * not use MHD_Connection::urh data anymore.
+   *
+   * In any mode #MHD_YES in this flag also means that
+   * MHD_Connection::urh data will not be used for socketpair
+   * forwarding and forwarding itself is finished.
+   *
+   * When BOTH @e was_closed (changed by command from application)
+   * AND @e clean_ready (changed internally by MHD) are set to
+   * #MHD_YES, function #MHD_resume_connection() will move this
+   * connection to cleanup list.
+   * @remark This flag could be changed from thread that process
+   * connection's recv(), send() and response.
+   */
+  int clean_ready;
 };
 
 
@@ -1410,11 +1439,15 @@ struct MHD_Daemon
 #if HTTPS_SUPPORT
   /**
    * Head of DLL of upgrade response handles we are processing.
+   * Used for upgraded TLS connections when thread-per-connection
+   * is not used.
    */
   struct MHD_UpgradeResponseHandle *urh_head;
 
   /**
    * Tail of DLL of upgrade response handles we are processing.
+   * Used for upgraded TLS connections when thread-per-connection
+   * is not used.
    */
   struct MHD_UpgradeResponseHandle *urh_tail;
 
diff --git a/src/microhttpd/response.c b/src/microhttpd/response.c
index b0fffe19..db77dc6d 100644
--- a/src/microhttpd/response.c
+++ b/src/microhttpd/response.c
@@ -21,6 +21,7 @@
  * @brief  Methods for managing response objects
  * @author Daniel Pittman
  * @author Christian Grothoff
+ * @author Karlson2k (Evgeny Grin)
  */
 
 #define MHD_NO_DEPRECATION 1
@@ -626,66 +627,41 @@ MHD_upgrade_action (struct MHD_UpgradeResponseHandle *urh,
                     enum MHD_UpgradeAction action,
                     ...)
 {
-  struct MHD_Connection *connection = urh->connection;
-  struct MHD_Daemon *daemon = connection->daemon;
+  struct MHD_Connection *connection;
+  struct MHD_Daemon *daemon;
+  if (NULL == urh)
+    return MHD_NO;
+  connection = urh->connection;
+
+  /* Precaution checks on external data. */
+  if (NULL == connection)
+    return MHD_NO;
+  daemon = connection->daemon;
+  if (NULL == daemon)
+    return MHD_NO;
 
   switch (action)
   {
   case MHD_UPGRADE_ACTION_CLOSE:
+    if (MHD_YES == urh->was_closed)
+      return MHD_NO; /* Already closed. */
+
     /* transition to special 'closed' state for start of cleanup */
+    urh->was_closed = MHD_YES;
     connection->state = MHD_CONNECTION_UPGRADE_CLOSED;
+    /* As soon as connection will be marked with BOTH
+     * 'urh->was_closed' AND 'urh->clean_ready', it will
+     * be moved to cleanup list by MHD_resume_connection(). */
+    MHD_resume_connection (connection);
 #if HTTPS_SUPPORT
     if (0 != (daemon->options & MHD_USE_TLS) )
       {
         /* signal that app is done by shutdown() of 'app' socket */
+        /* Application will not use anyway this socket after this command. */
         shutdown (urh->app.socket,
                   SHUT_RDWR);
       }
-#endif
-#if HTTPS_SUPPORT
-    if (0 != (daemon->options & MHD_USE_TLS) )
-      {
-        urh->was_closed = MHD_YES;
-        /* connection and urh cleanup will be done as soon as outgoing
-         * data will be sent and 'was_closed' is detected */
-        return MHD_YES;
-      }
-#endif
-    /* Application is done with this connection, tear it down! */
-    if (0 != (daemon->options & MHD_USE_THREAD_PER_CONNECTION) )
-      {
-        /* need to finish connection clean up */
-        MHD_cleanup_upgraded_connection_ (connection);
-        if (MHD_CONNECTION_IN_CLEANUP != connection->state)
-          {
-#if DEBUG_CLOSE
-#ifdef HAVE_MESSAGES
-            MHD_DLOG (connection->daemon,
-                      _("Processing thread terminating. Closing connection\n"));
-#endif
-#endif
-            if (MHD_CONNECTION_CLOSED != connection->state)
-              MHD_connection_close_ (connection,
-                                     MHD_REQUEST_TERMINATED_DAEMON_SHUTDOWN);
-            connection->idle_handler (connection);
-          }
-        if (NULL != connection->response)
-          {
-            MHD_destroy_response (connection->response);
-            connection->response = NULL;
-          }
-
-        if (MHD_INVALID_SOCKET != connection->socket_fd)
-          {
-            shutdown (connection->socket_fd,
-                      SHUT_WR);
-            MHD_socket_close_chk_ (connection->socket_fd);
-            connection->socket_fd = MHD_INVALID_SOCKET;
-          }
-        return MHD_YES;
-      }
-    /* 'upgraded' resources are not needed anymore - cleanup now */
-    MHD_cleanup_upgraded_connection_ (connection);
+#endif /* HTTPS_SUPPORT */
     return MHD_YES;
   default:
     /* we don't understand this one */
@@ -905,30 +881,25 @@ MHD_response_execute_upgrade_ (struct MHD_Response *response,
         DLL_insert (daemon->urh_head,
                     daemon->urh_tail,
                     urh);
-        /* Keep reference for later removal from the DLL */
-        connection->urh = urh;
       }
+    /* In thread-per-connection mode, thread will switch to forwarding once
+     * connection.urh is not NULL and connection.state == MHD_CONNECTION_UPGRADE.
+     */
   }
   else
     {
       urh->app.socket = MHD_INVALID_SOCKET;
       urh->mhd.socket = MHD_INVALID_SOCKET;
+      /* Non-TLS connection do not hold any additional resources. */
+      urh->clean_ready = MHD_YES;
     }
 #endif
-  if (0 != (daemon->options & MHD_USE_THREAD_PER_CONNECTION) )
-    {
-      /* Our caller will set 'connection->state' to
-         MHD_CONNECTION_UPGRADE, thereby triggering the main method
-         of the thread to switch to bi-directional forwarding or exit. */
-      connection->urh = urh;
-    }
-  else
-    {
-      /* As far as MHD's event loops are concerned, this connection is
-         suspended; it will be resumed once we are done in the
-         #MHD_upgrade_action() function */
-      MHD_suspend_connection (connection);
-    }
+  connection->urh = urh;
+  /* As far as MHD's event loops are concerned, this connection is
+     suspended; it will be resumed once application is done by the
+     #MHD_upgrade_action() function */
+  MHD_suspend_connection (connection);
+
   /* hand over socket to application */
   response->upgrade_handler (response->upgrade_handler_cls,
                              connection,
diff --git a/src/microhttpd/test_upgrade.c b/src/microhttpd/test_upgrade.c
index b369fd70..965ecc00 100644
--- a/src/microhttpd/test_upgrade.c
+++ b/src/microhttpd/test_upgrade.c
@@ -60,9 +60,8 @@ test_upgrade (int flags,
   struct sockaddr_in sa;
 
   done = 0;
-  if (0 == (flags & MHD_USE_THREAD_PER_CONNECTION))
-    flags |= MHD_USE_SUSPEND_RESUME;
-  d = MHD_start_daemon (flags | MHD_USE_DEBUG,
+
+  d = MHD_start_daemon (flags | MHD_USE_DEBUG | MHD_USE_SUSPEND_RESUME,
                         1080,
                         NULL, NULL,
                         &ahc_upgrade, NULL,
diff --git a/src/microhttpd/test_upgrade_ssl.c b/src/microhttpd/test_upgrade_ssl.c
index 07031cff..a72fad6d 100644
--- a/src/microhttpd/test_upgrade_ssl.c
+++ b/src/microhttpd/test_upgrade_ssl.c
@@ -138,9 +138,8 @@ test_upgrade (int flags,
   pid_t pid;
 
   done = 0;
-  if (0 == (flags & MHD_USE_THREAD_PER_CONNECTION))
-    flags |= MHD_USE_SUSPEND_RESUME;
-  d = MHD_start_daemon (flags | MHD_USE_DEBUG | MHD_USE_TLS,
+
+  d = MHD_start_daemon (flags | MHD_USE_DEBUG | MHD_USE_SUSPEND_RESUME |MHD_USE_TLS,
                         1080,
                         NULL, NULL,
                         &ahc_upgrade, NULL,