From 5e52950838956a5aa7e3e957753ea0c59dd2891a Mon Sep 17 00:00:00 2001 From: bel Date: Sun, 31 Aug 2014 01:34:04 +0200 Subject: [PATCH] Add close callback for websocket clients --- examples/websocket_client/websocket_client.c | 74 +++++++++++++++++--- include/civetweb.h | 10 ++- src/civetweb.c | 14 +++- test/unit_test.c | 20 +++--- 4 files changed, 92 insertions(+), 26 deletions(-) diff --git a/examples/websocket_client/websocket_client.c b/examples/websocket_client/websocket_client.c index 6359cf45..fe5638d7 100644 --- a/examples/websocket_client/websocket_client.c +++ b/examples/websocket_client/websocket_client.c @@ -78,7 +78,9 @@ struct mg_context * start_websocket_server() { const char * options[] = { "document_root", DOCUMENT_ROOT, "ssl_certificate", SSL_CERT, - "listening_ports", PORT, 0 + "listening_ports", PORT, + "request_timeout_ms", "5000", + 0 }; struct mg_callbacks callbacks; struct mg_context *ctx; @@ -100,6 +102,7 @@ struct mg_context * start_websocket_server() struct tclient_data { void * data; size_t len; + int closed; }; static int websocket_client_data_handler(struct mg_connection *conn, int flags, char *data, size_t data_len) @@ -107,7 +110,7 @@ static int websocket_client_data_handler(struct mg_connection *conn, int flags, struct mg_context *ctx = mg_get_context(conn); struct tclient_data *pclient_data = (struct tclient_data *) mg_get_user_data(ctx); - printf("From server: "); + printf("Client received data from server: "); fwrite(data, 1, data_len, stdout); printf("\n"); @@ -119,14 +122,25 @@ static int websocket_client_data_handler(struct mg_connection *conn, int flags, return 1; } +static void websocket_client_close_handler(struct mg_connection *conn) +{ + struct mg_context *ctx = mg_get_context(conn); + struct tclient_data *pclient_data = (struct tclient_data *) mg_get_user_data(ctx); + + printf("Client: Close handler\n"); + pclient_data->closed++; +} + int main(int argc, char *argv[]) { struct mg_context *ctx = NULL; - struct tclient_data client1_data = {NULL, 0}; - struct tclient_data client2_data = {NULL, 0}; + struct tclient_data client1_data = {NULL, 0, 0}; + struct tclient_data client2_data = {NULL, 0, 0}; + struct tclient_data client3_data = {NULL, 0, 0}; struct mg_connection* newconn1 = NULL; struct mg_connection* newconn2 = NULL; + struct mg_connection* newconn3 = NULL; char ebuf[100] = {0}; assert(websocket_welcome_msg_len == strlen(websocket_welcome_msg)); @@ -136,9 +150,10 @@ int main(int argc, char *argv[]) assert(ctx != NULL); printf("Server init\n\n"); - /* Then connect a client */ - newconn1 = mg_websocket_client_connect("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), - "/websocket", NULL, websocket_client_data_handler, &client1_data); + /* Then connect a first client */ + newconn1 = mg_connect_websocket_client("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), + "/websocket", NULL, websocket_client_data_handler, websocket_client_close_handler, + &client1_data); if (newconn1 == NULL) { @@ -147,6 +162,8 @@ int main(int argc, char *argv[]) } sleep(1); /* Should get the websocket welcome message */ + assert(client1_data.closed == 0); + assert(client2_data.closed == 0); assert(client2_data.data == NULL); assert(client2_data.len == 0); assert(client1_data.data != NULL); @@ -159,6 +176,8 @@ int main(int argc, char *argv[]) mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data1", 5); sleep(1); /* Should get the acknowledge message */ + assert(client1_data.closed == 0); + assert(client2_data.closed == 0); assert(client2_data.data == NULL); assert(client2_data.len == 0); assert(client1_data.data != NULL); @@ -168,9 +187,10 @@ int main(int argc, char *argv[]) client1_data.data = NULL; client1_data.len = 0; - /* Then connect a client */ - newconn2 = mg_websocket_client_connect("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), - "/websocket", NULL, websocket_client_data_handler, &client2_data); + /* Now connect a second client */ + newconn2 = mg_connect_websocket_client("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), + "/websocket", NULL, websocket_client_data_handler, websocket_client_close_handler, + &client2_data); if (newconn2 == NULL) { @@ -179,6 +199,8 @@ int main(int argc, char *argv[]) } sleep(1); /* Client 2 should get the websocket welcome message */ + assert(client1_data.closed == 0); + assert(client2_data.closed == 0); assert(client1_data.data == NULL); assert(client1_data.len == 0); assert(client2_data.data != NULL); @@ -191,6 +213,8 @@ int main(int argc, char *argv[]) mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data2", 5); sleep(1); /* Should get the acknowledge message */ + assert(client1_data.closed == 0); + assert(client2_data.closed == 0); assert(client2_data.data == NULL); assert(client2_data.len == 0); assert(client1_data.data != NULL); @@ -203,6 +227,8 @@ int main(int argc, char *argv[]) mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "bye", 3); sleep(1); /* Should get the goodbye message */ + assert(client1_data.closed == 0); + assert(client2_data.closed == 0); assert(client2_data.data == NULL); assert(client2_data.len == 0); assert(client1_data.data != NULL); @@ -215,6 +241,8 @@ int main(int argc, char *argv[]) mg_close_connection(newconn1); sleep(1); /* Won't get any message */ + assert(client1_data.closed == 1); + assert(client2_data.closed == 0); assert(client1_data.data == NULL); assert(client1_data.len == 0); assert(client2_data.data == NULL); @@ -223,6 +251,8 @@ int main(int argc, char *argv[]) mg_websocket_write(newconn2, WEBSOCKET_OPCODE_TEXT, "bye", 3); sleep(1); /* Should get the goodbye message */ + assert(client1_data.closed == 1); + assert(client2_data.closed == 0); assert(client1_data.data == NULL); assert(client1_data.len == 0); assert(client2_data.data != NULL); @@ -235,15 +265,39 @@ int main(int argc, char *argv[]) mg_close_connection(newconn2); sleep(1); /* Won't get any message */ + assert(client1_data.closed == 1); + assert(client2_data.closed == 1); assert(client1_data.data == NULL); assert(client1_data.len == 0); assert(client2_data.data == NULL); assert(client2_data.len == 0); + /* Connect client 3 */ + newconn3 = mg_connect_websocket_client("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), + "/websocket", NULL, websocket_client_data_handler, websocket_client_close_handler, + &client3_data); + + sleep(1); /* Client 3 should get the websocket welcome message */ + assert(client1_data.closed == 1); + assert(client2_data.closed == 1); + assert(client3_data.closed == 0); + assert(client1_data.data == NULL); + assert(client1_data.len == 0); + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + assert(client3_data.data != NULL); + assert(client3_data.len == websocket_welcome_msg_len); + assert(!memcmp(client3_data.data, websocket_welcome_msg, websocket_welcome_msg_len)); + free(client3_data.data); + client3_data.data = NULL; + client3_data.len = 0; + mg_stop(ctx); printf("Server shutdown\n"); sleep(10); + assert(client3_data.closed == 1); + return 0; } diff --git a/include/civetweb.h b/include/civetweb.h index 65a38911..27cfe562 100644 --- a/include/civetweb.h +++ b/include/civetweb.h @@ -582,6 +582,7 @@ CIVETWEB_API int mg_strncasecmp(const char *s1, const char *s2, size_t len); path: server path you are trying to connect to, i.e. if connection to localhost/app, path should be "/app" origin: value of the Origin HTTP header data_func: callback that should be used when data is received from the server + user_data: user supplied argument Return: On success, valid mg_connection object. @@ -589,11 +590,14 @@ CIVETWEB_API int mg_strncasecmp(const char *s1, const char *s2, size_t len); typedef int (*websocket_data_func)(struct mg_connection *, int bits, char *data, size_t data_len); -CIVETWEB_API struct mg_connection *mg_websocket_client_connect(const char *host, int port, int use_ssl, + +typedef void (*websocket_close_func)(struct mg_connection *); + +CIVETWEB_API struct mg_connection *mg_connect_websocket_client(const char *host, int port, int use_ssl, char *error_buffer, size_t error_buffer_size, const char *path, const char *origin, - websocket_data_func data_func, void * user_data); - + websocket_data_func data_func, websocket_close_func close_func, + void * user_data); #ifdef __cplusplus } #endif /* __cplusplus */ diff --git a/src/civetweb.c b/src/civetweb.c index 9e3904bf..2c079e28 100644 --- a/src/civetweb.c +++ b/src/civetweb.c @@ -6292,8 +6292,9 @@ static void close_connection(struct mg_connection *conn) #endif /* call the connection_close callback if assigned */ - if (conn->ctx->callbacks.connection_close != NULL) + if ((conn->ctx->callbacks.connection_close != NULL) && (conn->ctx->context_type == 1)) { conn->ctx->callbacks.connection_close(conn); + } mg_lock_connection(conn); @@ -6475,6 +6476,10 @@ static void* websocket_client_thread(void *data) DEBUG_TRACE("Websocket client thread exited\n"); + if (conn->ctx->callbacks.connection_close != NULL) { + conn->ctx->callbacks.connection_close(conn); + } + #ifdef _WIN32 return 0; #else @@ -6483,9 +6488,11 @@ static void* websocket_client_thread(void *data) } #endif -struct mg_connection *mg_websocket_client_connect(const char *host, int port, int use_ssl, +struct mg_connection *mg_connect_websocket_client(const char *host, int port, int use_ssl, char *error_buffer, size_t error_buffer_size, - const char *path, const char *origin, websocket_data_func data_func, void * user_data) + const char *path, const char *origin, + websocket_data_func data_func, websocket_close_func close_func, + void * user_data) { struct mg_connection* conn = NULL; struct mg_context * newctx = NULL; @@ -6534,6 +6541,7 @@ struct mg_connection *mg_websocket_client_connect(const char *host, int port, in newctx = (struct mg_context *) mg_malloc(sizeof(struct mg_context)); memcpy(newctx, conn->ctx, sizeof(struct mg_context)); newctx->callbacks.websocket_data = data_func; /* read_websocket will automatically call it */ + newctx->callbacks.connection_close = close_func; newctx->user_data = user_data; newctx->context_type = 2; /* client context type */ newctx->workerthreadcount = 1; /* one worker thread will be created */ diff --git a/test/unit_test.c b/test/unit_test.c index d2a305ec..05a9b3bc 100644 --- a/test/unit_test.c +++ b/test/unit_test.c @@ -541,15 +541,15 @@ static void test_mg_websocket_client_connect(int use_ssl) { /* Try to connect to our own server */ /* Invalid port test */ - conn = mg_websocket_client_connect("localhost", 0, use_ssl, + conn = mg_connect_websocket_client("localhost", 0, use_ssl, ebuf, sizeof(ebuf), - "/", "http://localhost", websocket_data_handler, NULL); + "/", "http://localhost", websocket_data_handler, NULL, NULL); ASSERT(conn == NULL); /* Should succeed, the default civetweb sever should complete the handshake */ - conn = mg_websocket_client_connect("localhost", port, use_ssl, + conn = mg_connect_websocket_client("localhost", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://localhost", websocket_data_handler, NULL); + "/", "http://localhost", websocket_data_handler, NULL, NULL); ASSERT(conn != NULL); /* Try an external server test */ @@ -557,21 +557,21 @@ static void test_mg_websocket_client_connect(int use_ssl) { if (use_ssl) { port = 443; } /* Not a websocket server path */ - conn = mg_websocket_client_connect("websocket.org", port, use_ssl, + conn = mg_connect_websocket_client("websocket.org", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org", websocket_data_handler, NULL); + "/", "http://websocket.org", websocket_data_handler, NULL, NULL); ASSERT(conn == NULL); /* Invalid port test */ - conn = mg_websocket_client_connect("echo.websocket.org", 0, use_ssl, + conn = mg_connect_websocket_client("echo.websocket.org", 0, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org", websocket_data_handler, NULL); + "/", "http://websocket.org", websocket_data_handler, NULL, NULL); ASSERT(conn == NULL); /* Should succeed, echo.websocket.org echos the data back */ - conn = mg_websocket_client_connect("echo.websocket.org", port, use_ssl, + conn = mg_connect_websocket_client("echo.websocket.org", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org", websocket_data_handler, NULL); + "/", "http://websocket.org", websocket_data_handler, NULL, NULL); ASSERT(conn != NULL); mg_stop(ctx);