Add close callback for websocket clients

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;
 }

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 */

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 */

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);