3 files changed, 48 insertions, 32 deletions

diff --git a/src/connection.c b/src/connection.c
index 58d86c31b..2015c9195 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -152,7 +152,7 @@ static void connSocketClose(connection *conn) {
     /* If called from within a handler, schedule the close but
      * keep the connection until the handler returns.
      */
-    if (conn->flags & CONN_FLAG_IN_HANDLER) {
+    if (connHasRefs(conn)) {
         conn->flags |= CONN_FLAG_CLOSE_SCHEDULED;
         return;
     }
@@ -183,10 +183,16 @@ static int connSocketRead(connection *conn, void *buf, size_t buf_len) {
 }
 
 static int connSocketAccept(connection *conn, ConnectionCallbackFunc accept_handler) {
+    int ret = C_OK;
+
     if (conn->state != CONN_STATE_ACCEPTING) return C_ERR;
     conn->state = CONN_STATE_CONNECTED;
-    if (!callHandler(conn, accept_handler)) return C_ERR;
-    return C_OK;
+
+    connIncrRefs(conn);
+    if (!callHandler(conn, accept_handler)) ret = C_ERR;
+    connDecrRefs(conn);
+
+    return ret;
 }
 
 /* Register a write handler, to be called when the connection is writable.
diff --git a/src/connection.h b/src/connection.h
index 97622f8d6..db09dfd83 100644
--- a/src/connection.h
+++ b/src/connection.h
@@ -45,9 +45,8 @@ typedef enum {
     CONN_STATE_ERROR
 } ConnectionState;
 
-#define CONN_FLAG_IN_HANDLER        (1<<0)      /* A handler execution is in progress */
-#define CONN_FLAG_CLOSE_SCHEDULED   (1<<1)      /* Closed scheduled by a handler */
-#define CONN_FLAG_WRITE_BARRIER     (1<<2)      /* Write barrier requested */
+#define CONN_FLAG_CLOSE_SCHEDULED   (1<<0)      /* Closed scheduled by a handler */
+#define CONN_FLAG_WRITE_BARRIER     (1<<1)      /* Write barrier requested */
 
 typedef void (*ConnectionCallbackFunc)(struct connection *conn);
 
@@ -70,7 +69,8 @@ typedef struct ConnectionType {
 struct connection {
     ConnectionType *type;
     ConnectionState state;
-    int flags;
+    short int flags;
+    short int refs;
     int last_errno;
     void *private_data;
     ConnectionCallbackFunc conn_handler;
@@ -88,6 +88,13 @@ struct connection {
  * connAccept() may directly call accept_handler(), or return and call it
  * at a later time. This behavior is a bit awkward but aims to reduce the need
  * to wait for the next event loop, if no additional handshake is required.
+ *
+ * IMPORTANT: accept_handler may decide to close the connection, calling connClose().
+ * To make this safe, the connection is only marked with CONN_FLAG_CLOSE_SCHEDULED
+ * in this case, and connAccept() returns with an error.
+ *
+ * connAccept() callers must always check the return value and on error (C_ERR)
+ * a connClose() must be called.
  */
 
 static inline int connAccept(connection *conn, ConnectionCallbackFunc accept_handler) {
diff --git a/src/connhelpers.h b/src/connhelpers.h
index f237c9b1d..86250d09e 100644
--- a/src/connhelpers.h
+++ b/src/connhelpers.h
@@ -37,46 +37,49 @@
  * implementations (currently sockets in connection.c and TLS in tls.c).
  *
  * Currently helpers implement the mechanisms for invoking connection
- * handlers, tracking in-handler states and dealing with deferred
- * destruction (if invoked by a handler).
+ * handlers and tracking connection references, to allow safe destruction
+ * of connections from within a handler.
  */
 
-/* Called whenever a handler is invoked on a connection and sets the
- * CONN_FLAG_IN_HANDLER flag to indicate we're in a handler context.
+/* Incremenet connection references.
  *
- * An attempt to close a connection while CONN_FLAG_IN_HANDLER is
- * set will result with deferred close, i.e. setting the CONN_FLAG_CLOSE_SCHEDULED
- * instead of destructing it.
+ * Inside a connection handler, we guarantee refs >= 1 so it is always
+ * safe to connClose().
+ *
+ * In other cases where we don't want to prematurely lose the connection,
+ * it can go beyond 1 as well; currently it is only done by connAccept().
  */
-static inline void enterHandler(connection *conn) {
-    conn->flags |= CONN_FLAG_IN_HANDLER;
+static inline void connIncrRefs(connection *conn) {
+    conn->refs++;
 }
 
-/* Called whenever a handler returns. This unsets the CONN_FLAG_IN_HANDLER
- * flag and performs actual close/destruction if a deferred close was
- * scheduled by the handler.
+/* Decrement connection references.
+ *
+ * Note that this is not intended to provide any automatic free logic!
+ * callHandler() takes care of that for the common flows, and anywhere an
+ * explicit connIncrRefs() is used, the caller is expected to take care of
+ * that.
  */
-static inline int exitHandler(connection *conn) {
-    conn->flags &= ~CONN_FLAG_IN_HANDLER;
-    if (conn->flags & CONN_FLAG_CLOSE_SCHEDULED) {
-        connClose(conn);
-        return 0;
-    }
-    return 1;
+
+static inline void connDecrRefs(connection *conn) {
+    conn->refs--;
+}
+
+static inline int connHasRefs(connection *conn) {
+    return conn->refs;
 }
 
 /* Helper for connection implementations to call handlers:
- * 1. Mark the handler in use.
+ * 1. Increment refs to protect the connection.
  * 2. Execute the handler (if set).
- * 3. Mark the handler as NOT in use and perform deferred close if was
- * requested by the handler at any time.
+ * 3. Decrement refs and perform deferred close, if refs==0.
  */
 static inline int callHandler(connection *conn, ConnectionCallbackFunc handler) {
-    conn->flags |= CONN_FLAG_IN_HANDLER;
+    connIncrRefs(conn);
     if (handler) handler(conn);
-    conn->flags &= ~CONN_FLAG_IN_HANDLER;
+    connDecrRefs(conn);
     if (conn->flags & CONN_FLAG_CLOSE_SCHEDULED) {
-        connClose(conn);
+        if (!connHasRefs(conn)) connClose(conn);
         return 0;
     }
     return 1;