4 files changed, 78 insertions, 67 deletions

diff --git a/src/microhttpd/daemon.c b/src/microhttpd/daemon.c
index 62cd3ce8..08709c80 100644
--- a/src/microhttpd/daemon.c
+++ b/src/microhttpd/daemon.c
@@ -883,6 +883,61 @@ call_handlers (struct MHD_Connection *con,
 
 #if HTTPS_SUPPORT
 /**
+ * This function finishes the process of closing the
+ * connection associated with the @a urh.  It should
+ * be called if the `was_closed` flag is set and the
+ * buffer has been drained.
+ *
+ * @param urh handle to the upgraded response we are finished with
+ */
+static void
+finish_upgrade_close (struct MHD_UpgradeResponseHandle *urh)
+{
+  struct MHD_Connection *connection = urh->connection;
+  struct MHD_Daemon *daemon = connection->daemon;
+
+  DLL_remove (daemon->urh_head,
+              daemon->urh_tail,
+              urh);
+  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");
+    }
+  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 ( (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");
+      if (0 != MHD_socket_close_ (urh->mhd.socket))
+        MHD_PANIC ("close failed\n");
+    }
+  MHD_resume_connection (connection);
+  MHD_connection_close_ (connection,
+                         MHD_REQUEST_TERMINATED_COMPLETED_OK);
+  free (urh);
+}
+
+
+/**
  * Performs bi-directional forwarding on upgraded HTTPS connections
  * based on the readyness state stored in the @a urh handle.
  *
@@ -891,6 +946,8 @@ call_handlers (struct MHD_Connection *con,
 static void
 process_urh (struct MHD_UpgradeResponseHandle *urh)
 {
+  int fin_read;
+
   /* handle reading from TLS client and writing to application */
   if ( (0 != (MHD_EPOLL_STATE_READ_READY & urh->app.celi)) &&
        (urh->in_buffer_off < urh->in_buffer_size) )
@@ -957,7 +1014,10 @@ process_urh (struct MHD_UpgradeResponseHandle *urh)
         {
           urh->out_buffer_off += res;
         }
+      fin_read = (0 == res);
     }
+  else
+    fin_read = 0;
   if ( (0 != (MHD_EPOLL_STATE_WRITE_READY & urh->app.celi)) &&
        (urh->out_buffer_off > 0) )
     {
@@ -986,6 +1046,10 @@ process_urh (struct MHD_UpgradeResponseHandle *urh)
             }
         }
     }
+  if ( (fin_read) &&
+       (0 == urh->out_buffer_off) &&
+       (MHD_YES == urh->was_closed) )
+    finish_upgrade_close (urh);
 }
 #endif
 
diff --git a/src/microhttpd/internal.h b/src/microhttpd/internal.h
index 7a1e58e1..9549b19a 100644
--- a/src/microhttpd/internal.h
+++ b/src/microhttpd/internal.h
@@ -1036,6 +1036,13 @@ struct MHD_UpgradeResponseHandle
    * nothing left.
    */
   char e_buf[RESERVE_EBUF_SIZE];
+
+  /**
+   * Set to #MHD_YES after the application closed the socket
+   * via #MHD_UPGRADE_ACTION_CLOSE.
+   */
+  int was_closed;
+
 #endif
 
 };
diff --git a/src/microhttpd/response.c b/src/microhttpd/response.c
index aa79196d..c08117d7 100644
--- a/src/microhttpd/response.c
+++ b/src/microhttpd/response.c
@@ -632,55 +632,27 @@ MHD_upgrade_action (struct MHD_UpgradeResponseHandle *urh,
             /* just need to signal the thread that we are done */
             MHD_semaphore_up (connection->upgrade_sem);
           }
+#if HTTPS_SUPPORT
         else
           {
             /* signal thread by shutdown() of 'app' socket */
-            shutdown (urh->app.socket, SHUT_RDWR);
+            shutdown (urh->app.socket,
+                      SHUT_RDWR);
           }
+#endif
         return MHD_YES;
       }
 #if HTTPS_SUPPORT
     if (0 != (daemon->options & MHD_USE_SSL) )
       {
-        DLL_remove (daemon->urh_head,
-                    daemon->urh_tail,
-                    urh);
-        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");
-          }
+        urh->was_closed = MHD_YES;
         if (MHD_INVALID_SOCKET != urh->app.socket)
           {
             if (0 != MHD_socket_close_ (urh->app.socket))
               MHD_PANIC ("close failed\n");
+            urh->app.socket = MHD_INVALID_SOCKET;
           }
-        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 ( (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");
-            if (0 != MHD_socket_close_ (urh->mhd.socket))
-              MHD_PANIC ("close failed\n");
-          }
+        return MHD_YES;
       }
 #endif
     MHD_resume_connection (connection);
@@ -688,34 +660,6 @@ MHD_upgrade_action (struct MHD_UpgradeResponseHandle *urh,
                            MHD_REQUEST_TERMINATED_COMPLETED_OK);
     free (urh);
     return MHD_YES;
-  case MHD_UPGRADE_ACTION_CORK:
-    /* FIXME: not implemented */
-    return MHD_NO;
-  case MHD_UPGRADE_ACTION_FLUSH:
-#if HTTPS_SUPPORT
-    if (0 != (daemon->options & MHD_USE_SSL))
-      {
-        int avail;
-
-        /* First, check that our pipe is empty, to be sure we do
-           have it all in the buffer. */
-        if ( (0 ==
-#if WINDOWS
-              ioctlsocket
-#else
-              ioctl
-#endif
-              (urh->mhd.socket,
-               FIONREAD,
-               &avail)) &&
-             (0 != avail) )
-          return MHD_NO;
-        /* then, refuse 'flush' unless our buffer is empty */
-        if (0 != urh->out_buffer_off)
-          return MHD_NO;
-      }
-#endif
-    return MHD_YES;
   default:
     /* we don't understand this one */
     return MHD_NO;
diff --git a/src/microhttpd/test_upgrade_ssl.c b/src/microhttpd/test_upgrade_ssl.c
index 5c663dd9..8077f402 100644
--- a/src/microhttpd/test_upgrade_ssl.c
+++ b/src/microhttpd/test_upgrade_ssl.c
@@ -259,10 +259,6 @@ run_usock (void *cls)
             "Finished");
   fprintf (stderr,
            "Closing socket\n");
-  while (MHD_NO ==
-         MHD_upgrade_action (urh,
-                             MHD_UPGRADE_ACTION_FLUSH))
-    usleep (1000);
   MHD_upgrade_action (urh,
                       MHD_UPGRADE_ACTION_CLOSE);
   fprintf (stderr,