From 9aaf36141e92ded84b6dd9abd3a08095f5593cdb Mon Sep 17 00:00:00 2001 From: "roberto@mrspurr" Date: Mon, 6 Dec 2010 13:48:23 +0100 Subject: [PATCH] extended the uwsgi api to allow easy implementation of websockets servers --- http.c | 113 ++++++++++++++++++++++++++------ plugins/python/uwsgi_pymodule.c | 54 +++++++++++++++ utils.c | 5 +- websockets.py | 48 ++++++++++++++ 4 files changed, 200 insertions(+), 20 deletions(-) create mode 100644 websockets.py diff --git a/http.c b/http.c index bbdd41bb..669072ea 100644 --- a/http.c +++ b/http.c @@ -176,10 +176,13 @@ static void *http_request(void *u_h_r) { int state = uwsgi_http_method; int http_body_len = 0; + int http_upgrade = 0; + + struct pollfd http_poll[2]; size_t len; - int i, j; + int i, j, rlen; char HTTP_header_key[1024]; @@ -279,6 +282,13 @@ static void *http_request(void *u_h_r) { *ptr++ = 0; http_body_len = atoi(tmp_buf); } + else if (!strcmp("CONNECTION", HTTP_header_key)) { + if (ptr+1 > watermark2) { close(uwsgi_fd); goto clear;} + *ptr++ = 0; + if (!strcmp(tmp_buf, "Upgrade")) { + http_upgrade = 1; + } + } ptr = tmp_buf; state = uwsgi_http_header_key; } else if (state == uwsgi_http_protocol_r) { @@ -322,31 +332,96 @@ static void *http_request(void *u_h_r) { uwsgi_error("write()"); } - if (http_body_len > 0) { - if (http_body_len >= (int) len - (i + 1)) { + if (http_upgrade) { + // send already available data + if ( (len - (i + 1)) > 0) { if (write(uwsgi_fd, buf + i + 1, len - (i + 1)) < 0) { - uwsgi_error("write()"); - } - http_body_len -= len - (i + 1); - } else { - if (write(uwsgi_fd, buf + i, http_body_len) < 0) { - uwsgi_error("write()"); - } - http_body_len = 0; + uwsgi_error("write()"); + close(uwsgi_fd); + goto clear; + } } - while (http_body_len > 0) { - int to_read = 4096; - if (http_body_len < to_read) { - to_read = http_body_len; + http_poll[0].fd = clientfd; + http_poll[0].events = POLLIN; + http_poll[1].fd = uwsgi_fd; + http_poll[1].events = POLLIN; + + for(;;) { + rlen = poll(http_poll, 2, -1); + if (rlen < 0) { + uwsgi_error("poll()"); + close(uwsgi_fd); + goto clear; } - len = read(clientfd, uwsgipkt, to_read); - if (write(uwsgi_fd, uwsgipkt, len) < 0) { - uwsgi_error("write()"); + else if (rlen > 0) { + if (http_poll[0].revents & POLLIN) { + len = read(clientfd, uwsgipkt, 4096); + if (len > 0) { + if (write(uwsgi_fd, uwsgipkt, len) < 0) { + uwsgi_error("write()"); + close(uwsgi_fd); + goto clear; + } + } + else { + // client disconnected + close(uwsgi_fd); + goto clear; + } + } + else if (http_poll[1].revents & POLLIN) { + len = read(uwsgi_fd, uwsgipkt, 4096); + if (len > 0) { + if (write(clientfd, uwsgipkt, len) < 0) { + uwsgi_error("write()"); + close(uwsgi_fd); + goto clear; + } + } + else { + // client disconnected + close(uwsgi_fd); + goto clear; + } + } + } + else { + // timeout + close(uwsgi_fd); + goto clear; + } + } + + } + else { + if (http_body_len > 0) { + if (http_body_len >= (int) len - (i + 1)) { + if (write(uwsgi_fd, buf + i + 1, len - (i + 1)) < 0) { + uwsgi_error("write()"); + } + http_body_len -= len - (i + 1); + } else { + if (write(uwsgi_fd, buf + i, http_body_len) < 0) { + uwsgi_error("write()"); + } + http_body_len = 0; + } + + while (http_body_len > 0) { + int to_read = 4096; + if (http_body_len < to_read) { + to_read = http_body_len; + } + len = read(clientfd, uwsgipkt, to_read); + if (write(uwsgi_fd, uwsgipkt, len) < 0) { + uwsgi_error("write()"); + } + http_body_len -= len; } - http_body_len -= len; } } + while ((len = read(uwsgi_fd, uwsgipkt, 4096)) > 0) { if (write(clientfd, uwsgipkt, len) < 0) { uwsgi_error("write()"); diff --git a/plugins/python/uwsgi_pymodule.c b/plugins/python/uwsgi_pymodule.c index 642c0491..04da9a0e 100644 --- a/plugins/python/uwsgi_pymodule.c +++ b/plugins/python/uwsgi_pymodule.c @@ -166,6 +166,59 @@ PyObject *py_uwsgi_close(PyObject * self, PyObject * args) { } +PyObject *py_uwsgi_recv_block(PyObject * self, PyObject * args) { + + char buf[4096]; + char *bufptr; + ssize_t rlen = 0, len ; + int fd, size, remains, ret, timeout = -1; + + + if (!PyArg_ParseTuple(args, "ii|i:recv_block", &fd, &size, &timeout)) { + return NULL; + } + + if (fd < 0) goto clear; + + UWSGI_RELEASE_GIL + // security check + if (size > 4096) size = 4096; + + remains = size; + + bufptr = buf; + while(remains > 0) { + uwsgi_log("%d %d %d\n", remains, size, timeout); + ret = uwsgi_waitfd(fd, timeout); + if (ret > 0) { + len = read(fd, bufptr, UMIN(remains, size)) ; + if (len > 0) { + bufptr+=len; + rlen += len; + remains -= len; + } + else { + break; + } + } + else { + uwsgi_log("error waiting for block data\n"); + break; + } + } + + UWSGI_GET_GIL + + if ( rlen == size) { + return PyString_FromStringAndSize(buf, rlen); + } + +clear: + + Py_INCREF(Py_None); + return Py_None; +} + PyObject *py_uwsgi_recv(PyObject * self, PyObject * args) { int fd, max_size = 4096; @@ -1639,6 +1692,7 @@ static PyMethodDef uwsgi_advanced_methods[] = { {"is_connected", py_uwsgi_is_connected, METH_VARARGS, ""}, {"send", py_uwsgi_send, METH_VARARGS, ""}, {"recv", py_uwsgi_recv, METH_VARARGS, ""}, + {"recv_block", py_uwsgi_recv_block, METH_VARARGS, ""}, {"close", py_uwsgi_close, METH_VARARGS, ""}, {"parsefile", py_uwsgi_parse_file, METH_VARARGS, ""}, diff --git a/utils.c b/utils.c index 7cd2cc9f..4789847d 100644 --- a/utils.c +++ b/utils.c @@ -1141,7 +1141,10 @@ int uwsgi_waitfd(int fd, int timeout) { if (!timeout) timeout = uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT]; - ret = poll(upoll, 1, timeout*1000); + timeout = timeout*1000; + if (timeout < 0) timeout = -1; + + ret = poll(upoll, 1, timeout); if (ret < 0) { uwsgi_error("poll()"); diff --git a/websockets.py b/websockets.py new file mode 100644 index 00000000..ec41d82e --- /dev/null +++ b/websockets.py @@ -0,0 +1,48 @@ +import uwsgi + +def application(e, s): + print e + + + client = e['wsgi.input'].fileno() + + print client + + data = uwsgi.recv_block(client, 8) + + print "data", data, len(data) + + key1 = e['HTTP_SEC_WEBSOCKET_KEY1'] + key2 = e['HTTP_SEC_WEBSOCKET_KEY2'] + + total1 = '' + div1 = 0 + for c in key1: + if c in '0'..'9': + total1 += c + + for c in key1: + if c == ' ': + div1 += 1 + + if div1 == 0: + raise StopIteration + + total1 = int(total1) / div1 + + total2 = '' + div2 = 0 + for c in key2: + if c in '0'..'9': + total2 += c + + for c in key2: + if c == ' ': + div2 += 1 + + if div2 == 0: + raise StopIteration + + total2 = int(total2) / div1 + +