From eac72465ad76a334c60b5a38f4b904c94a14d9e5 Mon Sep 17 00:00:00 2001 From: bel Date: Sat, 30 Aug 2014 23:31:21 +0200 Subject: [PATCH] Allow multiple websocket clients with different callbacks and an additional user supplied argument --- examples/websocket_client/websocket_client.c | 141 +++++++++++++------ include/civetweb.h | 7 +- src/civetweb.c | 53 +++++-- test/unit_test.c | 11 +- 4 files changed, 153 insertions(+), 59 deletions(-) diff --git a/examples/websocket_client/websocket_client.c b/examples/websocket_client/websocket_client.c index 14da7c0f..6359cf45 100644 --- a/examples/websocket_client/websocket_client.c +++ b/examples/websocket_client/websocket_client.c @@ -97,22 +97,24 @@ struct mg_context * start_websocket_server() /*************************************************************************************/ /* WEBSOCKET CLIENT */ /*************************************************************************************/ -struct { +struct tclient_data { void * data; size_t len; -} client_data; +}; static int websocket_client_data_handler(struct mg_connection *conn, int flags, char *data, size_t data_len) { + struct mg_context *ctx = mg_get_context(conn); + struct tclient_data *pclient_data = (struct tclient_data *) mg_get_user_data(ctx); + printf("From server: "); fwrite(data, 1, data_len, stdout); printf("\n"); - /* TODO: extra arg (instead of global client_data) */ - client_data.data = malloc(data_len); - assert(client_data.data != NULL); - memcpy(client_data.data, data, data_len); - client_data.len = data_len; + pclient_data->data = malloc(data_len); + assert(pclient_data->data != NULL); + memcpy(pclient_data->data, data, data_len); + pclient_data->len = data_len; return 1; } @@ -120,68 +122,123 @@ static int websocket_client_data_handler(struct mg_connection *conn, int flags, int main(int argc, char *argv[]) { - struct mg_context *ctx; - struct mg_connection* newconn; - char ebuf[100]; + struct mg_context *ctx = NULL; + struct tclient_data client1_data = {NULL, 0}; + struct tclient_data client2_data = {NULL, 0}; + struct mg_connection* newconn1 = NULL; + struct mg_connection* newconn2 = NULL; + char ebuf[100] = {0}; assert(websocket_welcome_msg_len == strlen(websocket_welcome_msg)); /* First set up a websocket server */ ctx = start_websocket_server(); assert(ctx != NULL); + printf("Server init\n\n"); /* Then connect a client */ - newconn = mg_websocket_client_connect("localhost", atoi(PORT), 0, - ebuf, sizeof(ebuf), - "/websocket", NULL, websocket_client_data_handler /* TODO: extra arg (instead of global client_data) */); + newconn1 = mg_websocket_client_connect("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), + "/websocket", NULL, websocket_client_data_handler, &client1_data); - if (newconn == NULL) + if (newconn1 == NULL) { printf("Error: %s", ebuf); return 1; } sleep(1); /* Should get the websocket welcome message */ - assert(client_data.data != NULL); - assert(client_data.len == websocket_welcome_msg_len); - assert(!memcmp(client_data.data, websocket_welcome_msg, websocket_welcome_msg_len)); - free(client_data.data); - client_data.data = NULL; - client_data.len = 0; + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + assert(client1_data.data != NULL); + assert(client1_data.len == websocket_welcome_msg_len); + assert(!memcmp(client1_data.data, websocket_welcome_msg, websocket_welcome_msg_len)); + free(client1_data.data); + client1_data.data = NULL; + client1_data.len = 0; - mg_websocket_write(newconn, WEBSOCKET_OPCODE_TEXT, "data1", 5); + mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data1", 5); sleep(1); /* Should get the acknowledge message */ - assert(client_data.data != NULL); - assert(client_data.len == websocket_acknowledge_msg_len); - assert(!memcmp(client_data.data, websocket_acknowledge_msg, websocket_acknowledge_msg_len)); - free(client_data.data); - client_data.data = NULL; - client_data.len = 0; + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + assert(client1_data.data != NULL); + assert(client1_data.len == websocket_acknowledge_msg_len); + assert(!memcmp(client1_data.data, websocket_acknowledge_msg, websocket_acknowledge_msg_len)); + free(client1_data.data); + client1_data.data = NULL; + client1_data.len = 0; - mg_websocket_write(newconn, WEBSOCKET_OPCODE_TEXT, "data2", 5); + /* Then connect a client */ + newconn2 = mg_websocket_client_connect("localhost", atoi(PORT), 0, ebuf, sizeof(ebuf), + "/websocket", NULL, websocket_client_data_handler, &client2_data); + + if (newconn2 == NULL) + { + printf("Error: %s", ebuf); + return 1; + } + + sleep(1); /* Client 2 should get the websocket welcome message */ + assert(client1_data.data == NULL); + assert(client1_data.len == 0); + assert(client2_data.data != NULL); + assert(client2_data.len == websocket_welcome_msg_len); + assert(!memcmp(client2_data.data, websocket_welcome_msg, websocket_welcome_msg_len)); + free(client2_data.data); + client2_data.data = NULL; + client2_data.len = 0; + + mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data2", 5); sleep(1); /* Should get the acknowledge message */ - assert(client_data.data != NULL); - assert(client_data.len == websocket_acknowledge_msg_len); - assert(!memcmp(client_data.data, websocket_acknowledge_msg, websocket_acknowledge_msg_len)); - free(client_data.data); - client_data.data = NULL; - client_data.len = 0; + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + assert(client1_data.data != NULL); + assert(client1_data.len == websocket_acknowledge_msg_len); + assert(!memcmp(client1_data.data, websocket_acknowledge_msg, websocket_acknowledge_msg_len)); + free(client1_data.data); + client1_data.data = NULL; + client1_data.len = 0; - mg_websocket_write(newconn, WEBSOCKET_OPCODE_TEXT, "bye", 3); + mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "bye", 3); sleep(1); /* Should get the goodbye message */ - assert(client_data.data != NULL); - assert(client_data.len == websocket_goodbye_msg_len); - assert(!memcmp(client_data.data, websocket_goodbye_msg, websocket_goodbye_msg_len)); - free(client_data.data); - client_data.data = NULL; - client_data.len = 0; + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + assert(client1_data.data != NULL); + assert(client1_data.len == websocket_goodbye_msg_len); + assert(!memcmp(client1_data.data, websocket_goodbye_msg, websocket_goodbye_msg_len)); + free(client1_data.data); + client1_data.data = NULL; + client1_data.len = 0; - mg_close_connection(newconn); + mg_close_connection(newconn1); sleep(1); /* Won't get any message */ + assert(client1_data.data == NULL); + assert(client1_data.len == 0); + assert(client2_data.data == NULL); + assert(client2_data.len == 0); + + mg_websocket_write(newconn2, WEBSOCKET_OPCODE_TEXT, "bye", 3); + + sleep(1); /* Should get the goodbye message */ + assert(client1_data.data == NULL); + assert(client1_data.len == 0); + assert(client2_data.data != NULL); + assert(client2_data.len == websocket_goodbye_msg_len); + assert(!memcmp(client2_data.data, websocket_goodbye_msg, websocket_goodbye_msg_len)); + free(client2_data.data); + client2_data.data = NULL; + client2_data.len = 0; + + mg_close_connection(newconn2); + + sleep(1); /* Won't get any message */ + assert(client1_data.data == NULL); + assert(client1_data.len == 0); + assert(client2_data.data == NULL); + assert(client2_data.len == 0); mg_stop(ctx); printf("Server shutdown\n"); diff --git a/include/civetweb.h b/include/civetweb.h index e4e3f39d..65a38911 100644 --- a/include/civetweb.h +++ b/include/civetweb.h @@ -573,7 +573,7 @@ CIVETWEB_API void mg_cry(struct mg_connection *conn, /* utility method to compare two buffers, case incensitive. */ CIVETWEB_API int mg_strncasecmp(const char *s1, const char *s2, size_t len); -/* Connect to a websocket as a client +/* Connect to a websocket as a client Parameters: host: host to connect to, i.e. "echo.websocket.org" or "192.168.1.1" or "localhost" port: server port @@ -586,12 +586,13 @@ CIVETWEB_API int mg_strncasecmp(const char *s1, const char *s2, size_t len); Return: On success, valid mg_connection object. On error, NULL. */ - + 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, char *error_buffer, size_t error_buffer_size, - const char *path, const char *origin, websocket_data_func data_func); + const char *path, const char *origin, + websocket_data_func data_func, void * user_data); #ifdef __cplusplus } diff --git a/src/civetweb.c b/src/civetweb.c index 4b64dbb9..9e3904bf 100644 --- a/src/civetweb.c +++ b/src/civetweb.c @@ -779,6 +779,7 @@ struct mg_context { char *config[NUM_OPTIONS]; /* Civetweb configuration parameters */ struct mg_callbacks callbacks; /* User-defined callback function */ void *user_data; /* User-defined data */ + int context_type; /* 1 = server context, 2 = client context */ struct socket *listening_sockets; in_port_t *listening_ports; @@ -5146,7 +5147,7 @@ static void read_websocket(struct mg_connection *conn) callback, and waiting repeatedly until an error occurs. */ /* TODO: Investigate if this next line is needed assert(conn->content_len == 0); */ - for (;;) { + while (!conn->ctx->stop_flag) { header_len = 0; assert(conn->data_len >= conn->request_len); if ((body_len = conn->data_len - conn->request_len) >= 2) { @@ -6316,12 +6317,29 @@ static void close_connection(struct mg_connection *conn) void mg_close_connection(struct mg_connection *conn) { + struct mg_context * client_ctx = NULL; + int i; + + if (conn->ctx->context_type == 2) { + client_ctx = conn->ctx; + /* client context: loops must end */ + conn->ctx->stop_flag = 1; + } + #ifndef NO_SSL if (conn->client_ssl_ctx != NULL) { SSL_CTX_free((SSL_CTX *) conn->client_ssl_ctx); } #endif close_connection(conn); + if (client_ctx != NULL) { + /* join worker thread and free context */ + for (i = 0; i < client_ctx->workerthreadcount; i++) { + mg_join_thread(client_ctx->workerthreadids[i]); + } + mg_free(client_ctx->workerthreadids); + mg_free(client_ctx); + } (void) pthread_mutex_destroy(&conn->mutex); mg_free(conn); } @@ -6446,26 +6464,35 @@ struct mg_connection *mg_download(const char *host, int port, int use_ssl, } #if defined(USE_WEBSOCKET) +#ifdef _WIN32 +static unsigned __stdcall websocket_client_thread(void *data) +#else static void* websocket_client_thread(void *data) +#endif { struct mg_connection* conn = (struct mg_connection*)data; read_websocket(conn); DEBUG_TRACE("Websocket client thread exited\n"); +#ifdef _WIN32 + return 0; +#else return NULL; +#endif } #endif struct mg_connection *mg_websocket_client_connect(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) + const char *path, const char *origin, websocket_data_func data_func, void * user_data) { struct mg_connection* conn = NULL; + struct mg_context * newctx = NULL; #if defined(USE_WEBSOCKET) static const char *magic = "x3JJHMbDL1EzLkh9GBhXDw=="; - static const char *handshake_req; + static const char *handshake_req; if(origin != NULL) { @@ -6502,15 +6529,24 @@ struct mg_connection *mg_websocket_client_connect(const char *host, int port, in return conn; } - /* For client connections, mg_context is fake. Set the callback for websocket - data manually here so that read_websocket will automatically call it */ - conn->ctx->callbacks.websocket_data = data_func; + /* For client connections, mg_context is fake. Since we need to set a callback + function, we need to create a copy and modify it. */ + 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->user_data = user_data; + newctx->context_type = 2; /* client context type */ + newctx->workerthreadcount = 1; /* one worker thread will be created */ + newctx->workerthreadids = (pthread_t*) mg_calloc(newctx->workerthreadcount, sizeof(pthread_t)); + conn->ctx = newctx; /* Start a thread to read the websocket client connection - This thread will automatically stop when mg_disconnect is + This thread will automatically stop when mg_disconnect is called on the client connection */ - if(mg_start_thread(websocket_client_thread, (void*)conn) != 0) + if (mg_start_thread_with_id(websocket_client_thread, (void*)conn, newctx->workerthreadids) != 0) { + mg_free((void*)newctx->workerthreadids); + mg_free((void*)newctx); mg_free((void*)conn); conn = NULL; DEBUG_TRACE("Websocket client connect thread could not be started\r\n"); @@ -7134,6 +7170,7 @@ struct mg_context *mg_start(const struct mg_callbacks *callbacks, ctx->callbacks.init_context(ctx); } ctx->callbacks.exit_context = exit_callback; + ctx->context_type = 1; /* server context */ /* Start master (listening) thread */ mg_start_thread_with_id(master_thread, ctx, &ctx->masterthreadid); diff --git a/test/unit_test.c b/test/unit_test.c index e9624f99..d2a305ec 100644 --- a/test/unit_test.c +++ b/test/unit_test.c @@ -543,16 +543,15 @@ static void test_mg_websocket_client_connect(int use_ssl) { /* Invalid port test */ conn = mg_websocket_client_connect("localhost", 0, use_ssl, ebuf, sizeof(ebuf), - "/", "http://localhost",websocket_data_handler); + "/", "http://localhost", websocket_data_handler, NULL); ASSERT(conn == NULL); /* Should succeed, the default civetweb sever should complete the handshake */ conn = mg_websocket_client_connect("localhost", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://localhost",websocket_data_handler); + "/", "http://localhost", websocket_data_handler, NULL); ASSERT(conn != NULL); - /* Try an external server test */ port = 80; if (use_ssl) { port = 443; } @@ -560,19 +559,19 @@ static void test_mg_websocket_client_connect(int use_ssl) { /* Not a websocket server path */ conn = mg_websocket_client_connect("websocket.org", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org",websocket_data_handler); + "/", "http://websocket.org", websocket_data_handler, NULL); ASSERT(conn == NULL); /* Invalid port test */ conn = mg_websocket_client_connect("echo.websocket.org", 0, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org",websocket_data_handler); + "/", "http://websocket.org", websocket_data_handler, NULL); ASSERT(conn == NULL); /* Should succeed, echo.websocket.org echos the data back */ conn = mg_websocket_client_connect("echo.websocket.org", port, use_ssl, ebuf, sizeof(ebuf), - "/", "http://websocket.org",websocket_data_handler); + "/", "http://websocket.org", websocket_data_handler, NULL); ASSERT(conn != NULL); mg_stop(ctx);