Browse code

modules/websocket: Closing handshake now working

- Also completed MI commands which allowed me to test Ping and Pong

Peter Dunkley authored on 21/06/2012 23:34:24
Showing 8 changed files
... ...
@@ -22,12 +22,23 @@
22 22
  */
23 23
 
24 24
 #include "../../locking.h"
25
+#include "../../str.h"
25 26
 #include "../../tcp_conn.h"
27
+#include "../../lib/kcore/kstats_wrapper.h"
28
+#include "../../lib/kmi/tree.h"
26 29
 #include "../../mem/mem.h"
27 30
 #include "ws_conn.h"
31
+#include "ws_mod.h"
32
+
33
+/* Maximum number of connections to display when using the ws.dump MI command */
34
+#define MAX_WS_CONNS_DUMP	50
28 35
 
29 36
 struct ws_connection **wsconn_hash = NULL;
30 37
 gen_lock_t *wsconn_lock = NULL;
38
+gen_lock_t *wsstat_lock = NULL;
39
+
40
+stat_var *ws_current_connections;
41
+stat_var *ws_max_concurrent_connections;
31 42
 
32 43
 char *wsconn_state_str[] =
33 44
 {
... ...
@@ -45,14 +56,24 @@ int wsconn_init(void)
45 56
 	if (wsconn_lock == NULL)
46 57
 	{
47 58
 		LM_ERR("allocating lock\n");
48
-		return -1;
59
+		goto error;
49 60
 	}
50 61
 	if (lock_init(wsconn_lock) == 0)
51 62
 	{
52 63
 		LM_ERR("initialising lock\n");
53
-		lock_dealloc((void *) wsconn_lock);
54
-		wsconn_lock = NULL;
55
-		return -1;
64
+		goto error;
65
+	}
66
+
67
+	wsstat_lock = lock_alloc();
68
+	if (wsstat_lock == NULL)
69
+	{
70
+		LM_ERR("allocating lock\n");
71
+		goto error;
72
+	}
73
+	if (lock_init(wsstat_lock) == NULL)
74
+	{
75
+		LM_ERR("initialising lock\n");
76
+		goto error;
56 77
 	}
57 78
 
58 79
 	wsconn_hash =
... ...
@@ -61,14 +82,19 @@ int wsconn_init(void)
61 82
 	if (wsconn_hash == NULL)
62 83
 	{
63 84
 		LM_ERR("allocating WebSocket hash-table\n");
64
-		lock_dealloc((void *) wsconn_lock);
65
-		wsconn_lock = NULL;
66
-		return -1;
85
+		goto error;
67 86
 	}
68 87
 	memset((void *) wsconn_hash, 0,
69 88
 		TCP_ID_HASH_SIZE * sizeof(ws_connection_t *));
70 89
 
71 90
 	return 0;
91
+
92
+error:
93
+	if (wsconn_lock) lock_dealloc((void *) wsconn_lock);
94
+	if (wsstat_lock) lock_dealloc((void *) wsstat_lock);
95
+	wsconn_lock = wsstat_lock = NULL;
96
+
97
+	return -1;
72 98
 }
73 99
 
74 100
 void wsconn_destroy(void)
... ...
@@ -77,8 +103,8 @@ void wsconn_destroy(void)
77 103
 
78 104
 	if (wsconn_hash)
79 105
 	{
80
-		WSCONN_UNLOCK;
81
-		WSCONN_LOCK;
106
+		lock_release(wsconn_lock);
107
+		lock_get(wsconn_lock);
82 108
 		for (h = 0; h < TCP_ID_HASH_SIZE; h++)
83 109
 		{
84 110
 			ws_connection_t *wsc = wsconn_hash[h];
... ...
@@ -89,7 +115,7 @@ void wsconn_destroy(void)
89 115
 				wsc = next;
90 116
 			}
91 117
 		}
92
-		WSCONN_UNLOCK;
118
+		lock_release(wsconn_lock);
93 119
 
94 120
 		shm_free(wsconn_hash);
95 121
 		wsconn_hash = NULL;
... ...
@@ -101,10 +127,18 @@ void wsconn_destroy(void)
101 127
 		lock_dealloc((void *) wsconn_lock);
102 128
 		wsconn_lock = NULL;
103 129
 	}
130
+
131
+	if (wsstat_lock)
132
+	{
133
+		lock_destroy(wsstat_lock);
134
+		lock_dealloc((void *) wsstat_lock);
135
+		wsstat_lock = NULL;
136
+	}
104 137
 }
105 138
 
106 139
 int wsconn_add(struct tcp_connection *con)
107 140
 {
141
+	int cur_cons, max_cons;
108 142
 	ws_connection_t *wsc;
109 143
 
110 144
 	if (!con)
... ...
@@ -130,12 +164,21 @@ int wsconn_add(struct tcp_connection *con)
130 164
 	   directly to this module */
131 165
 	con->flags |= F_CONN_WS;
132 166
 
133
-	WSCONN_LOCK;
167
+	lock_get(wsconn_lock);
134 168
 	wsc->next = wsconn_hash[wsc->id_hash];
135 169
 	wsc->prev = NULL;
136 170
 	if (wsconn_hash[wsc->id_hash]) wsconn_hash[wsc->id_hash]->prev = wsc;
137 171
 	wsconn_hash[wsc->id_hash] = wsc;
138
-	WSCONN_UNLOCK;
172
+	lock_release(wsconn_lock);
173
+
174
+	/* Update connection statistics */
175
+	lock_get(wsstat_lock);
176
+	update_stat(ws_current_connections, 1);
177
+	cur_cons = get_stat_val(ws_current_connections);
178
+	max_cons = get_stat_val(ws_max_concurrent_connections);
179
+	if (max_cons < cur_cons)
180
+		update_stat(ws_max_concurrent_connections, cur_cons - max_cons);
181
+	lock_release(wsstat_lock);
139 182
 
140 183
 	return 0;
141 184
 }
... ...
@@ -148,6 +191,7 @@ static inline void _wsconn_rm(ws_connection_t *wsc)
148 191
 	if (wsc->prev) wsc->prev->next = wsc->next;
149 192
 	shm_free(wsc);
150 193
 	wsc = NULL;
194
+	update_stat(ws_current_connections, -1);
151 195
 }
152 196
 
153 197
 int wsconn_rm(ws_connection_t *wsc)
... ...
@@ -158,13 +202,34 @@ int wsconn_rm(ws_connection_t *wsc)
158 202
 		return -1;
159 203
 	}
160 204
 
161
-	WSCONN_LOCK;
205
+	lock_get(wsconn_lock);
162 206
 	_wsconn_rm(wsc);
163
-	WSCONN_UNLOCK;
207
+	lock_release(wsconn_lock);
208
+
209
+	return 0;
210
+}
211
+
212
+int wsconn_update(ws_connection_t *wsc)
213
+{
214
+	if (!wsc)
215
+	{
216
+		LM_ERR("wsconn_rm: null pointer\n");
217
+		return -1;
218
+	}
164 219
 
220
+	wsc->last_used = (int) time(NULL);
165 221
 	return 0;
166 222
 }
167 223
 
224
+void wsconn_close_now(ws_connection_t *wsc)
225
+{
226
+	wsc->con->send_flags.f |= SND_F_CON_CLOSE;
227
+	wsc->con->state = S_CONN_BAD;
228
+	wsc->con->timeout = get_ticks_raw();
229
+	if (wsconn_rm(wsc) < 0)
230
+		LM_ERR("removing WebSocket connection\n");
231
+}
232
+
168 233
 ws_connection_t *wsconn_find(struct tcp_connection *con)
169 234
 {
170 235
 	ws_connection_t *wsc;
... ...
@@ -175,13 +240,85 @@ ws_connection_t *wsconn_find(struct tcp_connection *con)
175 240
 		return NULL;
176 241
 	}
177 242
 
178
-	WSCONN_LOCK;
243
+	lock_get(wsconn_lock);
179 244
 	for (wsc = wsconn_hash[con->id_hash]; wsc; wsc = wsc->next)
180 245
 	{
181
-		if (wsc->id_hash == con->id_hash)
246
+		if (wsc->con->id == con->id)
247
+		{
248
+			lock_release(wsconn_lock);
182 249
 			return wsc;
250
+		}
183 251
 	}
184
-	WSCONN_UNLOCK;
185 252
 
253
+	lock_release(wsconn_lock);
186 254
 	return NULL;
187 255
 }
256
+
257
+struct mi_root *ws_mi_dump(struct mi_root *cmd, void *param)
258
+{
259
+	int h, connections = 0, truncated = 0, interval;
260
+	char *src_proto, *dst_proto;
261
+	char src_ip[IP6_MAX_STR_SIZE + 1], dst_ip[IP6_MAX_STR_SIZE + 1];
262
+	ws_connection_t *wsc;
263
+	struct mi_root *rpl_tree = init_mi_tree(200, MI_OK_S, MI_OK_LEN);
264
+
265
+	if (!rpl_tree)
266
+		return 0;
267
+
268
+	lock_get(wsconn_lock);
269
+	for (h = 0; h < TCP_ID_HASH_SIZE; h++)
270
+	{
271
+		wsc = wsconn_hash[h];
272
+		while(wsc)
273
+		{
274
+			if (wsc->con)
275
+			{
276
+				src_proto = (wsc->con->rcv.proto== PROTO_TCP)
277
+						? "tcp" : "tls";
278
+				memset(src_ip, 0, IP6_MAX_STR_SIZE + 1);
279
+				ip_addr2sbuf(&wsc->con->rcv.src_ip, src_ip,
280
+						IP6_MAX_STR_SIZE);
281
+
282
+				dst_proto = (wsc->con->rcv.proto == PROTO_TCP)
283
+						? "tcp" : "tls";
284
+				memset(dst_ip, 0, IP6_MAX_STR_SIZE + 1);
285
+				ip_addr2sbuf(&wsc->con->rcv.dst_ip, src_ip,
286
+						IP6_MAX_STR_SIZE);
287
+
288
+				interval = (int)time(NULL) - wsc->last_used;
289
+
290
+				if (addf_mi_node_child(&rpl_tree->node, 0, 0, 0,
291
+						"%d: %s:%s:%hu -> %s:%s:%hu "
292
+						"(state: %s, "
293
+						"last used %ds ago)",
294
+						wsc->con->id,
295
+						src_proto,
296
+						strlen(src_ip) ? src_ip : "*",
297
+						wsc->con->rcv.src_port,
298
+						dst_proto,
299
+						strlen(dst_ip) ? dst_ip : "*",
300
+						wsc->con->rcv.dst_port,
301
+						wsconn_state_str[wsc->state],
302
+						interval) == 0)
303
+					return 0;
304
+
305
+				if (++connections == MAX_WS_CONNS_DUMP)
306
+				{
307
+					truncated = 1;
308
+					break;
309
+				}
310
+			}
311
+
312
+			wsc = wsc->next;
313
+		}
314
+	}
315
+	lock_release(wsconn_lock);
316
+
317
+	if (addf_mi_node_child(&rpl_tree->node, 0, 0, 0,
318
+				"%d WebSocket connection%s found%s",
319
+				connections, connections == 1 ? "" : "s",
320
+				truncated == 1 ? "(truncated)" : "") == 0)
321
+		return 0;
322
+
323
+	return rpl_tree;
324
+}
... ...
@@ -26,13 +26,14 @@
26 26
 
27 27
 #include "../../locking.h"
28 28
 #include "../../tcp_conn.h"
29
+#include "../../lib/kmi/tree.h"
29 30
 
30 31
 typedef enum
31 32
 {
32
-	WS_S_CONNECTING	= 0,
33
+	WS_S_CONNECTING	= 0,	/* Never used - included for completeness */
33 34
 	WS_S_OPEN,
34 35
 	WS_S_CLOSING,
35
-	WS_S_CLOSED
36
+	WS_S_CLOSED		/* Never used - included for completeness */
36 37
 } ws_conn_state_t;
37 38
 
38 39
 typedef struct ws_connection
... ...
@@ -40,24 +41,26 @@ typedef struct ws_connection
40 41
 	struct tcp_connection *con;
41 42
 
42 43
 	ws_conn_state_t state;
43
-	unsigned int id_hash;
44
-	unsigned int last_used;
44
+	int id;
45
+	unsigned id_hash;
46
+	int last_used;
45 47
 
46 48
 	struct ws_connection *prev;
47 49
 	struct ws_connection *next;
48 50
 } ws_connection_t;
49 51
 
50
-extern ws_connection_t **wsconn_hash;
51
-extern gen_lock_t *wsconn_lock;
52 52
 extern char *wsconn_state_str[];
53 53
 
54
+extern stat_var *ws_current_connections;
55
+extern stat_var *ws_max_concurrent_connections;
56
+
54 57
 int wsconn_init(void);
55 58
 void wsconn_destroy(void);
56 59
 int wsconn_add(struct tcp_connection *con);
57 60
 int wsconn_rm(ws_connection_t *wsc);
61
+int wsconn_update(ws_connection_t *wsc);
62
+void wsconn_close_now(ws_connection_t *wsc);
58 63
 ws_connection_t *wsconn_find(struct tcp_connection *con);
59
-
60
-#define WSCONN_LOCK	lock_get(wsconn_lock);
61
-#define WSCONN_UNLOCK	lock_release(wsconn_lock);
64
+struct mi_root *ws_mi_dump(struct mi_root *cmd, void *param);
62 65
 
63 66
 #endif /* _WS_CONN_H */
... ...
@@ -22,6 +22,7 @@
22 22
  */
23 23
 
24 24
 #include <limits.h>
25
+#include "../../str.h"
25 26
 #include "../../tcp_conn.h"
26 27
 #include "../../tcp_server.h"
27 28
 #include "../../lib/kcore/kstats_wrapper.h"
... ...
@@ -60,9 +61,15 @@ typedef struct {
60 61
 	unsigned int payload_len;
61 62
 	unsigned char masking_key[4];
62 63
 	char *payload_data;
63
-	tcp_event_info_t *tcpinfo;
64
+	ws_connection_t *wsc;
64 65
 } ws_frame_t;
65 66
 
67
+typedef enum
68
+{
69
+	CONN_CLOSE_DO = 0,
70
+	CONN_CLOSE_DONT
71
+} conn_close_t;
72
+
66 73
 #define BYTE0_MASK_FIN		(0x80)
67 74
 #define BYTE0_MASK_RSV1		(0x40)
68 75
 #define BYTE0_MASK_RSV2		(0x20)
... ...
@@ -80,19 +87,52 @@ typedef struct {
80 87
 #define OPCODE_PONG		(0xa)
81 88
 /* 0xb - 0xf are reserved for further control frames */
82 89
 
83
-
84
-static int decode_and_validate_ws_frame(ws_frame_t *frame)
90
+static int close_connection(ws_connection_t *wsc, ws_close_type_t type,
91
+				short int status, str reason);
92
+
93
+stat_var *ws_failed_connections;
94
+stat_var *ws_local_closed_connections;
95
+stat_var *ws_received_frames;
96
+stat_var *ws_remote_closed_connections;
97
+stat_var *ws_transmitted_frames;
98
+
99
+/* WebSocket status text */
100
+static str str_status_normal_closure = str_init("Normal closure");
101
+static str str_status_protocol_error = str_init("Protocol error");
102
+static str str_status_unsupported_opcode = str_init("Unsupported opcode");
103
+static str str_status_message_too_big = str_init("Message too big");
104
+
105
+/* MI command status text */
106
+static str str_status_empty_param = str_init("Empty connection ID parameter");
107
+static str str_status_too_many_params = str_init("Too many parameters");
108
+static str str_status_bad_param = str_init("Bad connection ID parameter");
109
+static str str_status_error_closing = str_init("Error closing connection");
110
+static str str_status_error_sending = str_init("Error sending frame");
111
+
112
+static int decode_and_validate_ws_frame(ws_frame_t *frame,
113
+					tcp_event_info_t *tcpinfo)
85 114
 {
86
-	unsigned int i, len=frame->tcpinfo->len;
115
+	unsigned int i, len = tcpinfo->len;
87 116
 	int mask_start, j;
88
-	char *buf = frame->tcpinfo->buf;
117
+	char *buf = tcpinfo->buf;
89 118
 
90 119
 	LM_INFO("decoding WebSocket frame\n");
91 120
 
121
+	if ((frame->wsc = wsconn_find(tcpinfo->con)) == NULL)
122
+	{
123
+		LM_WARN("WebSocket connection not found\n");
124
+		return -1;
125
+	}
126
+
127
+	wsconn_update(frame->wsc);
128
+
92 129
 	/* Decode and validate first 9 bits */
93 130
 	if (len < 2)
94 131
 	{
95 132
 		LM_WARN("message is too short\n");
133
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
134
+					str_status_protocol_error) < 0)
135
+			LM_ERR("closing connection\n");
96 136
 		return -1;
97 137
 	}
98 138
 	frame->fin = (buf[0] & 0xff) & BYTE0_MASK_FIN;
... ...
@@ -106,12 +146,18 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
106 146
 	{
107 147
 		LM_WARN("WebSocket fragmentation not supported in the sip "
108 148
 			"sub-protocol\n");
149
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
150
+					str_status_protocol_error) < 0)
151
+			LM_ERR("closing connection\n");
109 152
 		return -1;
110 153
 	}
111 154
 
112 155
 	if (frame->rsv1 || frame->rsv2 || frame->rsv3)
113 156
 	{
114 157
 		LM_WARN("WebSocket reserved fields with non-zero values\n");
158
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
159
+					str_status_protocol_error) < 0)
160
+			LM_ERR("closing connection\n");
115 161
 		return -1;
116 162
 	}
117 163
 
... ...
@@ -133,6 +179,9 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
133 179
 	default:
134 180
 		LM_WARN("unsupported opcode: 0x%x\n",
135 181
 			(unsigned char) frame->opcode);
182
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1008,
183
+					str_status_unsupported_opcode) < 0)
184
+			LM_ERR("closing connection\n");
136 185
 		return -1;
137 186
 	}
138 187
 
... ...
@@ -140,6 +189,9 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
140 189
 	{
141 190
 		LM_WARN("this is a server - all received messages must be "
142 191
 			"masked\n");
192
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
193
+					str_status_protocol_error) < 0)
194
+			LM_ERR("closing connection\n");
143 195
 		return -1;
144 196
 	}
145 197
 
... ...
@@ -150,6 +202,9 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
150 202
 		if (len < 4)
151 203
 		{
152 204
 			LM_WARN("message is too short\n");
205
+			if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
206
+						str_status_protocol_error) < 0)
207
+				LM_ERR("closing connection\n");
153 208
 			return -1;
154 209
 		}
155 210
 		mask_start = 4;
... ...
@@ -162,10 +217,23 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
162 217
 		if (len < 10)
163 218
 		{
164 219
 			LM_WARN("message is too short\n");
220
+			if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
221
+						str_status_protocol_error) < 0)
222
+				LM_ERR("closing connection\n");
165 223
 			return -1;
166 224
 		}
167 225
 		mask_start = 10;
168 226
 
227
+		if ((buf[2] & 0xff) != 0 || (buf[3] & 0xff) != 0
228
+			|| (buf[4] & 0xff) != 0 || (buf[5] & 0xff) != 0)
229
+		{
230
+			LM_WARN("message is too long\n");
231
+			if (close_connection(frame->wsc, LOCAL_CLOSE, 1009,
232
+						str_status_message_too_big) < 0)
233
+				LM_ERR("closing connection\n");
234
+			return -1;
235
+		}
236
+
169 237
 		/* Only decoding the last four bytes of the length...
170 238
 		   This limits the size of WebSocket messages that can be
171 239
 		   handled to 2^32 = which should be plenty for SIP! */
... ...
@@ -188,6 +256,9 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
188 256
 	{
189 257
 		LM_WARN("message not complete frame size %u but received %u\n",
190 258
 			frame->payload_len + mask_start + 4, len);
259
+		if (close_connection(frame->wsc, LOCAL_CLOSE, 1002,
260
+					str_status_protocol_error) < 0)
261
+			LM_ERR("closing connection\n");
191 262
 		return -1;
192 263
 	}
193 264
 	frame->payload_data = &buf[mask_start + 4];
... ...
@@ -204,7 +275,7 @@ static int decode_and_validate_ws_frame(ws_frame_t *frame)
204 275
 	return frame->opcode;
205 276
 }
206 277
 
207
-static int encode_and_send_ws_frame(ws_frame_t *frame, int conn_close)
278
+static int encode_and_send_ws_frame(ws_frame_t *frame, conn_close_t conn_close)
208 279
 {
209 280
 	int pos = 0, extended_length;
210 281
 	unsigned int frame_length;
... ...
@@ -213,6 +284,15 @@ static int encode_and_send_ws_frame(ws_frame_t *frame, int conn_close)
213 284
 
214 285
 	LM_INFO("encoding WebSocket frame\n");
215 286
 
287
+	if (frame->wsc->state != WS_S_OPEN)
288
+	{
289
+		LM_ERR("sending on closing connection\n");
290
+		wsconn_close_now(frame->wsc);
291
+		return -1;
292
+	}
293
+
294
+	wsconn_update(frame->wsc);
295
+
216 296
 	/* Validate the first byte */
217 297
 	if (!frame->fin)
218 298
 	{
... ...
@@ -292,23 +372,84 @@ static int encode_and_send_ws_frame(ws_frame_t *frame, int conn_close)
292 372
 	}
293 373
 	memcpy(&send_buf[pos], frame->payload_data, frame->payload_len);
294 374
 
295
-	init_dst_from_rcv(&dst, &frame->tcpinfo->con->rcv);
296
-	if (conn_close) dst.send_flags.f |= SND_F_CON_CLOSE;
375
+	init_dst_from_rcv(&dst, &frame->wsc->con->rcv);
376
+	if (conn_close == CONN_CLOSE_DO)
377
+	{
378
+		dst.send_flags.f |= SND_F_CON_CLOSE;
379
+		if (wsconn_rm(frame->wsc) < 0)
380
+		{
381
+			LM_ERR("removing WebSocket connection\n");
382
+			return -1;
383
+		}
384
+	}
297 385
 
298 386
 	if (tcp_send(&dst, NULL, send_buf, frame_length) < 0)
299 387
 	{
300 388
 		LM_ERR("sending WebSocket frame\n");
301 389
 		pkg_free(send_buf);
302 390
 		update_stat(ws_failed_connections, 1);
391
+		if (wsconn_rm(frame->wsc) < 0)
392
+			LM_ERR("removing WebSocket connection\n");
303 393
 		return -1;
304 394
 	}
305
-	
395
+
306 396
 	update_stat(ws_transmitted_frames, 1);
307 397
 
308 398
 	pkg_free(send_buf);
309 399
 	return 0;
310 400
 }
311 401
 
402
+static int close_connection(ws_connection_t *wsc, ws_close_type_t type,
403
+				short int status, str reason)
404
+{
405
+	char *data;
406
+	ws_frame_t frame;
407
+
408
+	data = pkg_malloc(sizeof(char) * (reason.len + 2));
409
+	if (data == NULL)
410
+	{
411
+		LM_ERR("allocating pkg memory\n");
412
+		return -1;
413
+	}
414
+
415
+	if (wsc->state == WS_S_OPEN)
416
+	{
417
+		data[0] = (status & 0xff00) >> 8;
418
+		data[1] = (status & 0x00ff) >> 0;
419
+		memcpy(&data[2], reason.s, reason.len);
420
+
421
+		memset(&frame, 0, sizeof(frame));
422
+		frame.fin = 1;
423
+		frame.opcode = OPCODE_CLOSE;
424
+		frame.payload_len = reason.len + 2;
425
+		frame.payload_data = data;
426
+		frame.wsc = wsc;
427
+
428
+		if (encode_and_send_ws_frame(&frame,
429
+			type ==
430
+			REMOTE_CLOSE ? CONN_CLOSE_DO : CONN_CLOSE_DONT) < 0)
431
+		{	
432
+			LM_ERR("sending WebSocket close\n");
433
+			pkg_free(data);
434
+			return -1;
435
+		}
436
+
437
+		pkg_free(data);
438
+
439
+		if (type == LOCAL_CLOSE)
440
+		{
441
+			frame.wsc->state = WS_S_CLOSING;
442
+			update_stat(ws_local_closed_connections, 1);
443
+		}
444
+		else
445
+			update_stat(ws_remote_closed_connections, 1);
446
+	}
447
+	else /* if (frame->wsc->state == WS_S_CLOSING) */
448
+		wsconn_close_now(wsc);
449
+
450
+	return 0;
451
+}
452
+
312 453
 static int handle_sip_message(ws_frame_t *frame)
313 454
 {
314 455
 	LM_INFO("Received SIP message\n");
... ...
@@ -323,10 +464,6 @@ static int handle_close(ws_frame_t *frame)
323 464
 	unsigned short code = 0;
324 465
 	str reason = {0, 0};
325 466
 
326
-	update_stat(ws_remote_closed_connections, 1);
327
-	update_stat(ws_current_connections, -1);
328
-	LM_INFO("Received Close\n");
329
-
330 467
 	if (frame->payload_len >= 2)
331 468
 		code =    ((frame->payload_data[0] & 0xff) << 8)
332 469
 			| ((frame->payload_data[1] & 0xff) << 0);
... ...
@@ -337,31 +474,33 @@ static int handle_close(ws_frame_t *frame)
337 474
 		reason.len = frame->payload_len - 2;
338 475
 	}
339 476
 
340
-	LM_INFO("Close: %hu %.*s\n", code, reason.len, reason.s); 
341
-
342
-	/* Close socket */
343
-	frame->tcpinfo->con->state = S_CONN_BAD;
344
-	frame->tcpinfo->con->timeout = get_ticks_raw();
477
+	LM_INFO("Received Close: %hu %.*s\n", code, reason.len, reason.s);
345 478
 
479
+	if (close_connection(frame->wsc,
480
+		frame->wsc->state == WS_S_OPEN ? REMOTE_CLOSE : LOCAL_CLOSE,
481
+		1000, str_status_normal_closure) < 0)
482
+	{
483
+		LM_ERR("closing connection\n");
484
+		return -1;
485
+	}
486
+	
346 487
 	return 0;
347 488
 }
348 489
 
349 490
 static int handle_ping(ws_frame_t *frame)
350 491
 {
351
-	LM_INFO("Received Ping\n");
492
+	LM_INFO("Received Ping: %.*s\n",
493
+		frame->payload_len, frame->payload_data);
352 494
 
353 495
 	frame->opcode = OPCODE_PONG;
354 496
 	frame->mask = 0;
355
-
356
-	encode_and_send_ws_frame(frame, 0);
497
+	encode_and_send_ws_frame(frame, CONN_CLOSE_DONT);
357 498
 
358 499
 	return 0;
359 500
 }
360 501
 
361 502
 static int handle_pong(ws_frame_t *frame)
362 503
 {
363
-	LM_INFO("Received Pong\n");
364
-
365 504
 	LM_INFO("Pong: %.*s\n", frame->payload_len, frame->payload_data);
366 505
 
367 506
 	return 0;
... ...
@@ -370,18 +509,17 @@ static int handle_pong(ws_frame_t *frame)
370 509
 int ws_frame_received(void *data)
371 510
 {
372 511
 	ws_frame_t ws_frame;
373
-	tcp_event_info_t *tev = (tcp_event_info_t *) data;
512
+	tcp_event_info_t *tcpinfo = (tcp_event_info_t *) data;
374 513
 
375 514
 	update_stat(ws_received_frames, 1);
376 515
 
377
-	if (tev == NULL || tev->buf == NULL || tev->len <= 0)
516
+	if (tcpinfo == NULL || tcpinfo->buf == NULL || tcpinfo->len <= 0)
378 517
 	{
379 518
 		LM_WARN("received bad frame\n");
380 519
 		return -1;
381 520
 	}
382 521
 
383
-	ws_frame.tcpinfo = tev;
384
-	switch(decode_and_validate_ws_frame(&ws_frame))
522
+	switch(decode_and_validate_ws_frame(&ws_frame, tcpinfo))
385 523
 	{
386 524
 	case OPCODE_TEXT_FRAME:
387 525
 	case OPCODE_BINARY_FRAME:
... ...
@@ -428,11 +566,7 @@ struct mi_root *ws_mi_close(struct mi_root *cmd, void *param)
428 566
 {
429 567
 	unsigned int id;
430 568
 	struct mi_node *node = NULL;
431
-	ws_frame_t frame;
432
-	tcp_event_info_t tcpinfo;
433
-	short int code = 1000;
434
-	str reason = str_init("Normal Closure");
435
-	char *data;
569
+	ws_connection_t *wsc;
436 570
 
437 571
 	node = cmd->node.kids;
438 572
 	if (node == NULL)
... ...
@@ -440,7 +574,8 @@ struct mi_root *ws_mi_close(struct mi_root *cmd, void *param)
440 574
 	if (node->value.s == NULL || node->value.len == 0)
441 575
 	{
442 576
 		LM_ERR("empty connection ID parameter\n");
443
-		return init_mi_tree(400, "Empty connection ID parameter", 29);
577
+		return init_mi_tree(400, str_status_empty_param.s,
578
+					str_status_empty_param.len);
444 579
 	}
445 580
 	if (str2int(&node->value, &id) < 0)
446 581
 	{
... ...
@@ -450,54 +585,98 @@ struct mi_root *ws_mi_close(struct mi_root *cmd, void *param)
450 585
 	if (node->next != NULL)
451 586
 	{
452 587
 		LM_ERR("too many parameters\n");
453
-		return init_mi_tree(400, "Too many parameters", 19);
588
+		return init_mi_tree(400, str_status_too_many_params.s,
589
+					str_status_too_many_params.len);
454 590
 	}
455 591
 
456
-	if ((tcpinfo.con = tcpconn_get(id, 0, 0, 0, 0)) == NULL)
592
+	if ((wsc = wsconn_find(tcpconn_get(id, 0, 0, 0, 0))) == NULL)
457 593
 	{
458 594
 		LM_ERR("bad connection ID parameter\n");
459
-		return init_mi_tree(400, "Bad connection ID parameter", 27);
595
+		return init_mi_tree(400, str_status_bad_param.s,
596
+					str_status_bad_param.len);
460 597
 	}
461 598
 
462
-	if ((data = pkg_malloc(sizeof(char) * (reason.len + 2))) == NULL)
599
+	if (close_connection(wsc, LOCAL_CLOSE, 1000,
600
+				str_status_normal_closure) < 0)
463 601
 	{
464
-		LM_ERR("allocating pkg memory\n");
465
-		return 0;
602
+		LM_ERR("closing connection\n");
603
+		return init_mi_tree(500, str_status_error_closing.s,
604
+					str_status_error_closing.len);
466 605
 	}
467 606
 
468
-	data[0] = (code & 0xff00) >> 8;
469
-	data[1] = (code & 0x00ff) >> 0;
470
-	memcpy(&data[2], reason.s, reason.len);
607
+	return init_mi_tree(200, MI_OK_S, MI_OK_LEN);
608
+}
609
+
610
+static int ping_pong(ws_connection_t *wsc, int opcode)
611
+{
612
+	ws_frame_t frame;
471 613
 
472 614
 	memset(&frame, 0, sizeof(frame));
473 615
 	frame.fin = 1;
474
-	frame.opcode = OPCODE_CLOSE;
475
-	frame.payload_len = reason.len + 2;
476
-	frame.payload_data = data;
477
-	frame.tcpinfo = &tcpinfo;
616
+	frame.opcode = opcode;
617
+	frame.payload_len = server_hdr.len;
618
+	frame.payload_data = server_hdr.s;
619
+	frame.wsc = wsc;
620
+
621
+	if (encode_and_send_ws_frame(&frame, CONN_CLOSE_DONT) < 0)
622
+	{	
623
+		LM_ERR("closing connection\n");
624
+		return -1;
625
+	}
626
+
627
+	return 0;
628
+}
629
+
630
+static struct mi_root *mi_ping_pong(struct mi_root *cmd, void *param,
631
+					int opcode)
632
+{
633
+	unsigned int id;
634
+	struct mi_node *node = NULL;
635
+	ws_connection_t *wsc;
478 636
 
479
-	if (encode_and_send_ws_frame(&frame, 1) < 0)
637
+	node = cmd->node.kids;
638
+	if (node == NULL)
639
+		return 0;
640
+	if (node->value.s == NULL || node->value.len == 0)
480 641
 	{
481
-		LM_ERR("sending WebSocket close\n");
482
-		pkg_free(data);
483
-		return init_mi_tree(500,"Sending WebSocket close", 23);
642
+		LM_ERR("empty connection ID parameter\n");
643
+		return init_mi_tree(400, str_status_empty_param.s,
644
+					str_status_empty_param.len);
645
+	}
646
+	if (str2int(&node->value, &id) < 0)
647
+	{
648
+		LM_ERR("converting string to int\n");
649
+		return 0;
650
+	}
651
+	if (node->next != NULL)
652
+	{
653
+		LM_ERR("too many parameters\n");
654
+		return init_mi_tree(400, str_status_too_many_params.s,
655
+					str_status_too_many_params.len);
484 656
 	}
485 657
 
486
-	update_stat(ws_local_closed_connections, 1);
487
-	update_stat(ws_current_connections, -1);
658
+	if ((wsc = wsconn_find(tcpconn_get(id, 0, 0, 0, 0))) == NULL)
659
+	{
660
+		LM_ERR("bad connection ID parameter\n");
661
+		return init_mi_tree(400, str_status_bad_param.s,
662
+					str_status_bad_param.len);
663
+	}
664
+
665
+	if (ping_pong(wsc, opcode) < 0)
666
+	{
667
+		return init_mi_tree(500, str_status_error_sending.s,
668
+					str_status_error_sending.len);
669
+	}
488 670
 
489
-	pkg_free(data);
490 671
 	return init_mi_tree(200, MI_OK_S, MI_OK_LEN);
491 672
 }
492 673
 
493 674
 struct mi_root *ws_mi_ping(struct mi_root *cmd, void *param)
494 675
 {
495
-	/* TODO Ping specified connection */
496
-	return init_mi_tree(200, MI_OK_S, MI_OK_LEN);
676
+	return mi_ping_pong(cmd, param, OPCODE_PING);
497 677
 }
498 678
 
499 679
 struct mi_root *ws_mi_pong(struct mi_root *cmd, void *param)
500 680
 {
501
-	/* TODO Pong specified connection */
502
-	return init_mi_tree(200, MI_OK_S, MI_OK_LEN);
681
+	return mi_ping_pong(cmd, param, OPCODE_PONG);
503 682
 }
... ...
@@ -25,7 +25,21 @@
25 25
 #define _WS_FRAME_H
26 26
 
27 27
 #include "../../sr_module.h"
28
+#include "../../str.h"
28 29
 #include "../../lib/kmi/tree.h"
30
+#include "ws_conn.h"
31
+
32
+typedef enum
33
+{
34
+	LOCAL_CLOSE = 0,
35
+	REMOTE_CLOSE
36
+} ws_close_type_t;
37
+
38
+extern stat_var *ws_failed_connections;
39
+extern stat_var *ws_local_closed_connections;
40
+extern stat_var *ws_received_frames;
41
+extern stat_var *ws_remote_closed_connections;
42
+extern stat_var *ws_transmitted_frames;
29 43
 
30 44
 int ws_frame_received(void *data);
31 45
 struct mi_root *ws_mi_close(struct mi_root *cmd, void *param);
... ...
@@ -27,6 +27,7 @@
27 27
 #include "../../data_lump_rpl.h"
28 28
 #include "../../dprint.h"
29 29
 #include "../../locking.h"
30
+#include "../../str.h"
30 31
 #include "../../tcp_conn.h"
31 32
 #include "../../lib/kcore/kstats_wrapper.h"
32 33
 #include "../../lib/kcore/cmpapi.h"
... ...
@@ -41,6 +42,9 @@
41 42
 
42 43
 #define WS_VERSION		(13)
43 44
 
45
+stat_var *ws_failed_handshakes;
46
+stat_var *ws_successful_handshakes;
47
+
44 48
 static str str_sip = str_init("sip");
45 49
 static str str_upgrade = str_init("upgrade");
46 50
 static str str_websocket = str_init("websocket");
... ...
@@ -79,8 +83,6 @@ static char key_buf[KEY_BUF_LEN];
79 83
 
80 84
 static int ws_send_reply(sip_msg_t *msg, int code, str *reason, str *hdrs)
81 85
 {
82
-	int cur_cons, max_cons;
83
-
84 86
 	if (hdrs && hdrs->len > 0)
85 87
 	{
86 88
 		if (add_lump_rpl(msg, hdrs->s, hdrs->len, LUMP_RPL_HDR) == 0)
... ...
@@ -98,23 +100,9 @@ static int ws_send_reply(sip_msg_t *msg, int code, str *reason, str *hdrs)
98 100
 		return -1;
99 101
 	}
100 102
 
101
-	if (code == 101)
102
-	{
103
-		update_stat(ws_successful_handshakes, 1);
104
-
105
-		lock_get(ws_stats_lock);
106
-		update_stat(ws_current_connections, 1);
107
-
108
-		cur_cons = get_stat_val(ws_current_connections);
109
-		max_cons = get_stat_val(ws_max_concurrent_connections);
110
-
111
-		if (max_cons < cur_cons)
112
-			update_stat(ws_max_concurrent_connections,
113
-						cur_cons - max_cons);
114
-		lock_release(ws_stats_lock);
115
-	}
116
-	else
117
-		update_stat(ws_failed_handshakes, 1);
103
+	update_stat(
104
+		code == 101 ? ws_successful_handshakes : ws_failed_handshakes,
105
+		1);
118 106
 
119 107
 	return 0;
120 108
 }
... ...
@@ -27,6 +27,9 @@
27 27
 #include "../../sr_module.h"
28 28
 #include "../../parser/msg_parser.h"
29 29
 
30
+stat_var *ws_failed_handshakes;
31
+stat_var *ws_successful_handshakes;
32
+
30 33
 int ws_handle_handshake(struct sip_msg *msg);
31 34
 struct mi_root *ws_mi_disable(struct mi_root *cmd, void *param);
32 35
 struct mi_root *ws_mi_enable(struct mi_root *cmd, void *param);
... ...
@@ -29,7 +29,6 @@
29 29
 #include "../../tcp_conn.h"
30 30
 #include "../../lib/kcore/kstats_wrapper.h"
31 31
 #include "../../lib/kmi/mi.h"
32
-#include "../../lib/kmi/tree.h"
33 32
 #include "../../mem/mem.h"
34 33
 #include "../../parser/msg_parser.h"
35 34
 #include "ws_conn.h"
... ...
@@ -42,29 +41,13 @@ MODULE_VERSION
42 41
 /* Maximum number of connections to display when using the ws.dump MI command */
43 42
 #define MAX_WS_CONNS_DUMP	50
44 43
 
45
-extern gen_lock_t *tcpconn_lock;
46
-extern struct tcp_connection **tcpconn_id_hash;
47
-
48 44
 static int mod_init(void);
49 45
 static void destroy(void);
50 46
 
51 47
 sl_api_t ws_slb;
52 48
 int *ws_enabled;
53
-gen_lock_t *ws_stats_lock;
54
-
55
-int ws_ping_interval = 30;	/* time (in seconds) between sending Pings */
56
-
57
-stat_var *ws_current_connections;
58
-stat_var *ws_failed_connections;
59
-stat_var *ws_failed_handshakes;
60
-stat_var *ws_local_closed_connections;
61
-stat_var *ws_max_concurrent_connections;
62
-stat_var *ws_received_frames;
63
-stat_var *ws_remote_closed_connections;
64
-stat_var *ws_successful_handshakes;
65
-stat_var *ws_transmitted_frames;
66 49
 
67
-static struct mi_root *mi_dump(struct mi_root *cmd, void *param);
50
+int ws_ping_interval = 180;	/* time (in seconds) between sending Pings */
68 51
 
69 52
 static cmd_export_t cmds[]= 
70 53
 {
... ...
@@ -77,19 +60,24 @@ static cmd_export_t cmds[]=
77 60
 static param_export_t params[]=
78 61
 {
79 62
 	{ "ping_interval",	INT_PARAM, &ws_ping_interval },
80
-	{ 0, 0 }
63
+	{ 0, 0, 0 }
81 64
 };
82 65
 
83 66
 static stat_export_t stats[] =
84 67
 {
68
+	/* ws_conn.c */
85 69
 	{ "ws_current_connections",       0, &ws_current_connections },
86
-	{ "ws_failed_connections",        0, &ws_failed_connections },
70
+	{ "ws_max_concurrent_connections",0, &ws_max_concurrent_connections },
71
+
72
+	/* ws_handshake.c */
87 73
 	{ "ws_failed_handshakes",         0, &ws_failed_handshakes },
74
+	{ "ws_successful_handshakes",     0, &ws_successful_handshakes },
75
+
76
+	/* ws_frame.c */
77
+	{ "ws_failed_connections",        0, &ws_failed_connections },
88 78
 	{ "ws_local_closed_connections",  0, &ws_local_closed_connections },
89
-	{ "ws_max_concurrent_connections",0, &ws_max_concurrent_connections },
90 79
 	{ "ws_received_frames",           0, &ws_received_frames },
91 80
 	{ "ws_remote_closed_connections", 0, &ws_remote_closed_connections },
92
-	{ "ws_successful_handshakes",     0, &ws_successful_handshakes },
93 81
 	{ "ws_transmitted_frames",        0, &ws_transmitted_frames },
94 82
 	{ 0, 0, 0 }
95 83
 };
... ...
@@ -98,7 +86,7 @@ static mi_export_t mi_cmds[] =
98 86
 {
99 87
 	{ "ws.close",   ws_mi_close,   0, 0, 0 },
100 88
 	{ "ws.disable", ws_mi_disable, 0, 0, 0 },
101
-	{ "ws.dump",	mi_dump,       0, 0, 0 },
89
+	{ "ws.dump",	ws_mi_dump,    0, 0, 0 },
102 90
 	{ "ws.enable",	ws_mi_enable,  0, 0, 0 },
103 91
 	{ "ws.ping",    ws_mi_ping,    0, 0, 0 },
104 92
 	{ "ws.pong",	ws_mi_pong,    0, 0, 0 },
... ...
@@ -160,31 +148,10 @@ static int mod_init(void)
160 148
 	}
161 149
 	*ws_enabled = 1;
162 150
 
163
-	if (wsconn_init() < 0)
164
-	{
165
-		LM_ERR("initialising WebSocket connections table\n");
166
-		goto error;
167
-	}
168
-
169
-	if ((ws_stats_lock = lock_alloc()) == NULL)
170
-	{
171
-		LM_ERR("allocating lock\n");
172
-		goto error;
173
-	}
174
-	if (lock_init(ws_stats_lock) == NULL)
175
-	{
176
-		LM_ERR("initialising lock\n");
177
-		goto error;
178
-	}
179
-
180 151
 	return 0;
181 152
 
182 153
 error:
183 154
 	wsconn_destroy();
184
-
185
-	if (ws_stats_lock)
186
-		lock_dealloc(ws_stats_lock);
187
-
188 155
 	shm_free(ws_enabled);
189 156
 
190 157
 	return -1;
... ...
@@ -194,75 +161,4 @@ static void destroy(void)
194 161
 {
195 162
 	wsconn_destroy();
196 163
 	shm_free(ws_enabled);
197
-	lock_destroy(ws_stats_lock);
198
-	lock_dealloc(ws_stats_lock);
199
-}
200
-
201
-static struct mi_root *mi_dump(struct mi_root *cmd, void *param)
202
-{
203
-	int h, connections = 0, truncated = 0, interval;
204
-	char *src_proto, *dst_proto;
205
-	char src_ip[IP6_MAX_STR_SIZE + 1], dst_ip[IP6_MAX_STR_SIZE + 1];
206
-	ws_connection_t *wsc;
207
-	struct mi_root *rpl_tree = init_mi_tree(200, MI_OK_S, MI_OK_LEN);
208
-
209
-	if (!rpl_tree)
210
-		return 0;
211
-
212
-	WSCONN_LOCK;
213
-	for (h = 0; h < TCP_ID_HASH_SIZE; h++)
214
-	{
215
-		wsc = wsconn_hash[h];
216
-		while(wsc)
217
-		{
218
-			if (wsc->con)
219
-			{
220
-				src_proto = (wsc->con->rcv.proto== PROTO_TCP)
221
-						? "tcp" : "tls";
222
-				memset(src_ip, 0, IP6_MAX_STR_SIZE + 1);
223
-				ip_addr2sbuf(&wsc->con->rcv.src_ip, src_ip,
224
-						IP6_MAX_STR_SIZE);
225
-
226
-				dst_proto = (wsc->con->rcv.proto == PROTO_TCP)
227
-						? "tcp" : "tls";
228
-				memset(dst_ip, 0, IP6_MAX_STR_SIZE + 1);
229
-				ip_addr2sbuf(&wsc->con->rcv.dst_ip, src_ip,
230
-						IP6_MAX_STR_SIZE);
231
-
232
-				interval = (int)time(NULL) - wsc->last_used;
233
-
234
-				if (addf_mi_node_child(&rpl_tree->node, 0, 0, 0,
235
-						"%d: %s:%s:%hu -> %s:%s:%hu "
236
-						"(state: %s, "
237
-						"last used %ds ago)",
238
-						wsc->con->id,
239
-						src_proto,
240
-						strlen(src_ip) ? src_ip : "*",
241
-						wsc->con->rcv.src_port,
242
-						dst_proto,
243
-						strlen(dst_ip) ? dst_ip : "*",
244
-						wsc->con->rcv.dst_port,
245
-						wsconn_state_str[wsc->state],
246
-						interval) == 0)
247
-					return 0;
248
-
249
-				if (++connections == MAX_WS_CONNS_DUMP)
250
-				{
251
-					truncated = 1;
252
-					break;
253
-				}
254
-			}
255
-
256
-			wsc = wsc->next;
257
-		}
258
-	}
259
-	WSCONN_UNLOCK;
260
-
261
-	if (addf_mi_node_child(&rpl_tree->node, 0, 0, 0,
262
-				"%d WebSocket connection%s found%s",
263
-				connections, connections == 1 ? "" : "s",
264
-				truncated == 1 ? "(truncated)" : "") == 0)
265
-		return 0;
266
-
267
-	return rpl_tree;
268 164
 }
... ...
@@ -34,14 +34,4 @@ extern gen_lock_t *ws_stats_lock;
34 34
 
35 35
 extern int ws_ping_interval;	/* time (in seconds) between sending Pings */
36 36
 
37
-extern stat_var *ws_current_connections;
38
-extern stat_var *ws_failed_connections;
39
-extern stat_var *ws_failed_handshakes;
40
-extern stat_var *ws_local_closed_connections;
41
-extern stat_var *ws_max_concurrent_connections;
42
-extern stat_var *ws_received_frames;
43
-extern stat_var *ws_remote_closed_connections;
44
-extern stat_var *ws_successful_handshakes;
45
-extern stat_var *ws_transmitted_frames;
46
-
47 37
 #endif /* _WS_MOD_H */