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;