diff --git a/lib/handshake.c b/lib/handshake.c index 41ba01c6..4cc698e5 100644 --- a/lib/handshake.c +++ b/lib/handshake.c @@ -91,17 +91,8 @@ handshake_00(struct libwebsocket_context *context, struct libwebsocket *wsi) goto bail; /* allocate the per-connection user memory (if any) */ - - if (wsi->protocol->per_session_data_size) { - wsi->user_space = malloc( - wsi->protocol->per_session_data_size); - if (wsi->user_space == NULL) { - fprintf(stderr, "Out of memory for " - "conn user space\n"); - goto bail; - } - } else - wsi->user_space = NULL; + if (wsi->protocol->per_session_data_size && !libwebsocket_ensure_user_space(wsi)) + goto bail; /* create the response packet */ @@ -275,17 +266,8 @@ handshake_0405(struct libwebsocket_context *context, struct libwebsocket *wsi) } /* allocate the per-connection user memory (if any) */ - - if (wsi->protocol->per_session_data_size) { - wsi->user_space = malloc( - wsi->protocol->per_session_data_size); - if (wsi->user_space == NULL) { - fprintf(stderr, "Out of memory for " - "conn user space\n"); - goto bail; - } - } else - wsi->user_space = NULL; + if (wsi->protocol->per_session_data_size && !libwebsocket_ensure_user_space(wsi)) + goto bail; /* create the response packet */ diff --git a/lib/libwebsockets.c b/lib/libwebsockets.c index 50aeaee2..7d691c49 100644 --- a/lib/libwebsockets.c +++ b/lib/libwebsockets.c @@ -1373,17 +1373,8 @@ check_accept: accept_ok: /* allocate the per-connection user memory (if any) */ - - if (wsi->protocol->per_session_data_size) { - wsi->user_space = malloc( - wsi->protocol->per_session_data_size); - if (wsi->user_space == NULL) { - fprintf(stderr, "Out of memory for " - "conn user space\n"); - goto bail2; - } - } else - wsi->user_space = NULL; + if (wsi->protocol->per_session_data_size && !libwebsocket_ensure_user_space(wsi)) + goto bail2; /* clear his proxy connection timeout */ @@ -3096,3 +3087,21 @@ libwebsocket_is_final_fragment(struct libwebsocket *wsi) { return wsi->final; } + +void * +libwebsocket_ensure_user_space(struct libwebsocket *wsi) +{ + /* allocate the per-connection user memory (if any) */ + + if (wsi->protocol->per_session_data_size && !wsi->user_space) { + wsi->user_space = malloc( + wsi->protocol->per_session_data_size); + if (wsi->user_space == NULL) { + fprintf(stderr, "Out of memory for " + "conn user space\n"); + return NULL; + } + memset(wsi->user_space, 0, wsi->protocol->per_session_data_size); + } + return wsi->user_space; +} diff --git a/lib/libwebsockets.h b/lib/libwebsockets.h index bb414286..d65d8c18 100644 --- a/lib/libwebsockets.h +++ b/lib/libwebsockets.h @@ -647,6 +647,9 @@ libwebsocket_get_socket_fd(struct libwebsocket *wsi); LWS_EXTERN int libwebsocket_is_final_fragment(struct libwebsocket *wsi); +LWS_EXTERN void * +libwebsocket_ensure_user_space(struct libwebsocket *wsi); + LWS_EXTERN int libwebsocket_rx_flow_control(struct libwebsocket *wsi, int enable);