/tests/websocket.c

  1/* http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76 */
  2
  3#include <stdlib.h>
  4#define _GNU_SOURCE
  5#include <stdio.h>
  6#include <string.h>
  7#include <unistd.h>
  8#include <pthread.h>
  9#include <arpa/inet.h>
 10#include <errno.h>
 11
 12#include <sys/types.h>
 13#include <sys/socket.h>
 14
 15#include <event.h>
 16
 17struct host_info {
 18	char *host;
 19	short port;
 20};
 21
 22/* worker_thread, with counter of remaining messages */
 23struct worker_thread {
 24	struct host_info *hi;
 25	struct event_base *base;
 26
 27	int msg_target;
 28	int msg_received;
 29	int msg_sent;
 30	int byte_count;
 31	pthread_t thread;
 32
 33	struct evbuffer *buffer;
 34	int got_header;
 35
 36	int verbose;
 37	struct event ev_w;
 38};
 39
 40void
 41process_message(struct worker_thread *wt, size_t sz) {
 42
 43	// printf("process_message\n");
 44	if(wt->msg_received % 10000 == 0) {
 45		printf("thread %u: %8d messages left (got %9d bytes so far).\n",
 46			(unsigned int)wt->thread,
 47			wt->msg_target - wt->msg_received, wt->byte_count);
 48	}
 49	wt->byte_count += sz;
 50
 51	/* decrement read count, and stop receiving when we reach zero. */
 52	wt->msg_received++;
 53	if(wt->msg_received == wt->msg_target) {
 54		event_base_loopexit(wt->base, NULL);
 55	}
 56}
 57
 58void
 59websocket_write(int fd, short event, void *ptr) {
 60	int ret;
 61	struct worker_thread *wt = ptr;
 62
 63	if(event != EV_WRITE) {
 64		return;
 65	}
 66
 67	char message[] = "\x00[\"SET\",\"key\",\"value\"]\xff\x00[\"GET\",\"key\"]\xff";
 68	ret = write(fd, message, sizeof(message)-1);
 69	if(ret != sizeof(message)-1) {
 70		fprintf(stderr, "write on %d failed: %s\n", fd, strerror(errno));
 71		close(fd);
 72	}
 73
 74	wt->msg_sent += 2;
 75	if(wt->msg_sent < wt->msg_target) {
 76		event_set(&wt->ev_w, fd, EV_WRITE, websocket_write, wt);
 77		event_base_set(wt->base, &wt->ev_w);
 78		ret = event_add(&wt->ev_w, NULL);
 79	}
 80}
 81
 82static void
 83websocket_read(int fd, short event, void *ptr) {
 84	char packet[2048], *pos;
 85	int ret, success = 1;
 86
 87	struct worker_thread *wt = ptr;
 88
 89	if(event != EV_READ) {
 90		return;
 91	}
 92
 93	/* read message */
 94	ret = read(fd, packet, sizeof(packet));
 95	pos = packet;
 96	if(ret > 0) {
 97		char *data, *last;
 98		int sz, msg_sz;
 99
100		if(wt->got_header == 0) { /* first response */
101			char *frame_start = strstr(packet, "MH"); /* end of the handshake */
102			if(frame_start == NULL) {
103				return; /* not yet */
104			} else { /* start monitoring possible writes */
105				printf("start monitoring possible writes\n");
106				evbuffer_add(wt->buffer, frame_start + 2, ret - (frame_start + 2 - packet));
107
108				wt->got_header = 1;
109				event_set(&wt->ev_w, fd, EV_WRITE,
110						websocket_write, wt);
111				event_base_set(wt->base, &wt->ev_w);
112				ret = event_add(&wt->ev_w, NULL);
113			}
114		} else {
115			/* we've had the header already, now bufffer data. */
116			evbuffer_add(wt->buffer, packet, ret);
117		}
118
119		while(1) {
120			data = (char*)EVBUFFER_DATA(wt->buffer);
121			sz = EVBUFFER_LENGTH(wt->buffer);
122
123			if(sz == 0) { /* no data */
124				break;
125			}
126			if(*data != 0) { /* missing frame start */
127				success = 0;
128				break;
129			}
130			last = memchr(data, 0xff, sz); /* look for frame end */
131			if(!last) {
132				/* no end of frame in sight. */
133				break;
134			}
135			msg_sz = last - data - 1;
136			process_message(ptr, msg_sz); /* record packet */
137
138			/* drain including frame delimiters (+2 bytes) */
139			evbuffer_drain(wt->buffer, msg_sz + 2); 
140		}
141	} else {
142		printf("ret=%d\n", ret);
143		success = 0;
144	}
145	if(success == 0) {
146		shutdown(fd, SHUT_RDWR);
147		close(fd);
148		event_base_loopexit(wt->base, NULL);
149	}
150}
151
152void*
153worker_main(void *ptr) {
154
155	char ws_template[] = "GET /.json HTTP/1.1\r\n"
156				"Host: %s:%d\r\n"
157				"Connection: Upgrade\r\n"
158				"Upgrade: WebSocket\r\n"
159				"Origin: http://%s:%d\r\n"
160				"Sec-WebSocket-Key1: 18x 6]8vM;54 *(5:  {   U1]8  z [  8\r\n"
161				"Sec-WebSocket-Key2: 1_ tx7X d  <  nw  334J702) 7]o}` 0\r\n"
162				"\r\n"
163				"Tm[K T2u";
164
165	struct worker_thread *wt = ptr;
166
167	int ret;
168	int fd;
169	struct sockaddr_in addr;
170	char *ws_handshake;
171	size_t ws_handshake_sz;
172
173	/* connect socket */
174	fd = socket(AF_INET, SOCK_STREAM, 0);
175	addr.sin_family = AF_INET;
176	addr.sin_port = htons(wt->hi->port);
177	memset(&(addr.sin_addr), 0, sizeof(addr.sin_addr));
178	addr.sin_addr.s_addr = inet_addr(wt->hi->host);
179
180	ret = connect(fd, (struct sockaddr*)&addr, sizeof(struct sockaddr));
181	if(ret != 0) {
182		fprintf(stderr, "connect: ret=%d: %s\n", ret, strerror(errno));
183		return NULL;
184	}
185
186	/* initialize worker thread */
187	wt->base = event_base_new();
188	wt->buffer = evbuffer_new();
189	wt->byte_count = 0;
190	wt->got_header = 0;
191
192	/* send handshake */
193	ws_handshake_sz = sizeof(ws_handshake)
194		+ 2*strlen(wt->hi->host) + 500;
195	ws_handshake = calloc(ws_handshake_sz, 1);
196	ws_handshake_sz = (size_t)sprintf(ws_handshake, ws_template, 
197			wt->hi->host, wt->hi->port,
198			wt->hi->host, wt->hi->port);
199	ret = write(fd, ws_handshake, ws_handshake_sz);
200
201	struct event ev_r;
202	event_set(&ev_r, fd, EV_READ | EV_PERSIST, websocket_read, wt);
203	event_base_set(wt->base, &ev_r);
204	event_add(&ev_r, NULL);
205
206	/* go! */
207	event_base_dispatch(wt->base);
208	event_base_free(wt->base);
209	free(ws_handshake);
210	return NULL;
211}
212
213void
214usage(const char* argv0, char *host_default, short port_default,
215		int thread_count_default, int messages_default) {
216
217	printf("Usage: %s [options]\n"
218		"Options are:\n"
219		"\t-h host\t\t(default = \"%s\")\n"
220		"\t-p port\t\t(default = %d)\n"
221		"\t-c threads\t(default = %d)\n"
222		"\t-n count\t(number of messages per thread, default = %d)\n"
223		"\t-v\t\t(verbose)\n",
224		argv0, host_default, (int)port_default,
225		thread_count_default, messages_default);
226}
227
228int
229main(int argc, char *argv[]) {
230
231	struct timespec t0, t1;
232
233	int messages_default = 100000;
234	int thread_count_default = 4;
235	short port_default = 7379;
236	char *host_default = "127.0.0.1";
237
238	int msg_target = messages_default;
239	int thread_count = thread_count_default;
240	int i, opt;
241	char *colon;
242	double total = 0, total_bytes = 0;
243	int verbose = 0;
244
245	struct host_info hi = {host_default, port_default};
246
247	struct worker_thread *workers;
248	
249	/* getopt */
250	while ((opt = getopt(argc, argv, "h:p:c:n:v")) != -1) {
251		switch (opt) {
252			case 'h':
253				colon = strchr(optarg, ':');
254				if(!colon) {
255					size_t sz = strlen(optarg);
256					hi.host = calloc(1 + sz, 1);
257					strncpy(hi.host, optarg, sz);
258				} else {
259					hi.host = calloc(1+colon-optarg, 1);
260					strncpy(hi.host, optarg, colon-optarg);
261					hi.port = (short)atol(colon+1);
262				}
263				break;
264
265			case 'p':
266				hi.port = (short)atol(optarg);
267				break;
268
269			case 'c':
270				thread_count = atoi(optarg);
271				break;
272
273			case 'n':
274				msg_target = atoi(optarg);
275				break;
276
277			case 'v':
278				verbose = 1;
279				break;
280			default: /* '?' */
281				usage(argv[0], host_default, port_default,
282						thread_count_default,
283						messages_default);
284				exit(EXIT_FAILURE);
285		}
286	}
287
288	/* run threads */
289	workers = calloc(sizeof(struct worker_thread), thread_count);
290
291	clock_gettime(CLOCK_MONOTONIC, &t0);
292	for(i = 0; i < thread_count; ++i) {
293		workers[i].msg_target = msg_target;
294		workers[i].hi = &hi;
295		workers[i].verbose = verbose;
296
297		pthread_create(&workers[i].thread, NULL,
298				worker_main, &workers[i]);
299	}
300
301	/* wait for threads to finish */
302	for(i = 0; i < thread_count; ++i) {
303		pthread_join(workers[i].thread, NULL);
304		total += workers[i].msg_received;
305		total_bytes += workers[i].byte_count;
306	}
307
308	/* timing */
309	clock_gettime(CLOCK_MONOTONIC, &t1);
310	float mili0 = t0.tv_sec * 1000 + t0.tv_nsec / 1000000;
311	float mili1 = t1.tv_sec * 1000 + t1.tv_nsec / 1000000;
312
313	if(total != 0) {
314		printf("Read %ld messages in %0.2f sec: %0.2f msg/sec (%d MB/sec, %d KB/sec)\n",
315			(long)total, 
316			(mili1-mili0)/1000.0,
317			1000*total/(mili1-mili0),
318			(int)(total_bytes / (1000*(mili1-mili0))),
319			(int)(total_bytes / (mili1-mili0)));
320		return EXIT_SUCCESS;
321	} else {
322		printf("No message was read.\n");
323		return EXIT_FAILURE;
324	}
325}
326