Create new websocket function websocket_check_pingpong

core/websockets.c

@@ -216,7 +216,7 @@ static ssize_t uwsgi_websockets_recv_pkt(struct wsgi_request *wsgi_req, int nb)
 }
 
 
-static struct uwsgi_buffer *uwsgi_websocket_recv_do(struct wsgi_request *wsgi_req, int nb) {
+struct uwsgi_buffer *uwsgi_websocket_recv_do(struct wsgi_request *wsgi_req, int nb) {
 	if (!wsgi_req->websocket_buf) {
 		// this buffer will be destroyed on connection close
 		wsgi_req->websocket_buf = uwsgi_buffer_new(uwsgi.page_size);

plugins/python/uwsgi_pymodule.c

@@ -1276,7 +1276,7 @@ PyObject *py_uwsgi_websocket_recv_nb(PyObject * self, PyObject * args, PyObject
         struct uwsgi_buffer *ub = uwsgi_websocket_recv_nb(wsgi_req);
         UWSGI_GET_GIL
         if (!ub) {
-                return PyErr_Format(PyExc_IOError, "unable to receive websocket message");
+            return PyErr_Format(PyExc_IOError, "unable to receive websocket message");
         }
 
         PyObject *ret = PyString_FromStringAndSize(ub->buf, ub->pos);
@@ -1285,6 +1285,41 @@ PyObject *py_uwsgi_websocket_recv_nb(PyObject * self, PyObject * args, PyObject
 }
 
 
+PyObject *py_uwsgi_websocket_check_pingpong(PyObject * self, PyObject *args, PyObject *kwargs) {
+	PyObject *request_context = NULL;
+
+	static char *kwlist[] = {"request_context", NULL};
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O:websocket_recv_nb", kwlist, &request_context)) {
+		return NULL;
+	}
+
+	struct wsgi_request *wsgi_req;
+	if(request_context == NULL) {
+		wsgi_req = py_current_wsgi_req();
+	} else {
+		py_uwsgi_context_argument_check(request_context);
+		wsgi_req = py_current_wsgi_req_from_context(request_context);
+	}
+
+	if(wsgi_req == NULL) {
+		return PyErr_Format(PyExc_IOError, "Unable to send ping/pong");
+	}
+
+	UWSGI_RELEASE_GIL
+	struct uwsgi_buffer *ub = uwsgi_websocket_recv_do(wsgi_req, 1);  // Non blocking by default
+	UWSGI_GET_GIL
+	if (!ub) {
+		wsgi_req->websocket_closed = 1;
+		return PyErr_Format(PyExc_IOError, "Pong not received");
+	}
+
+	PyObject *ret = PyString_FromStringAndSize(ub->buf, ub->pos);
+        uwsgi_buffer_destroy(ub);
+        return ret;
+}
+
+
 PyObject *py_uwsgi_embedded_data(PyObject * self, PyObject * args) {
 
 	char *name;
@@ -2867,6 +2902,7 @@ static PyMethodDef uwsgi_advanced_methods[] = {
 	{"websocket_send", (PyCFunction)py_uwsgi_websocket_send, METH_VARARGS|METH_KEYWORDS|METH_KEYWORDS, ""},
 	{"websocket_send_binary", (PyCFunction)py_uwsgi_websocket_send_binary, METH_VARARGS|METH_KEYWORDS, ""},
 	{"websocket_handshake", py_uwsgi_websocket_handshake, METH_VARARGS, ""},
+	{"websocket_check_pingpong", (PyCFunction)py_uwsgi_websocket_check_pingpong, METH_VARARGS|METH_KEYWORDS, ""},
 
 	{"chunked_read", py_uwsgi_chunked_read, METH_VARARGS, ""},
 	{"chunked_read_nb", py_uwsgi_chunked_read_nb, METH_VARARGS, ""},

uwsgi.h

@@ -4620,8 +4620,10 @@ void uwsgi_subscribe_all(uint8_t, int);
 void uwsgi_websockets_init(void);
 int uwsgi_websocket_send(struct wsgi_request *, char *, size_t);
 int uwsgi_websocket_send_binary(struct wsgi_request *, char *, size_t);
+int uwsgi_websocket_check_pingpong(struct wsgi_request *);
 struct uwsgi_buffer *uwsgi_websocket_recv(struct wsgi_request *);
 struct uwsgi_buffer *uwsgi_websocket_recv_nb(struct wsgi_request *);
+struct uwsgi_buffer *uwsgi_websocket_recv_do(struct wsgi_request *, int);
 
 char *uwsgi_chunked_read(struct wsgi_request *, size_t *, int, int);
 struct uwsgi_buffer *uwsgi_chunked_read_smart(struct wsgi_request *, size_t, int);