improve websocket implementation

README.md

@@ -12,6 +12,8 @@ graphical interface served on the webpage.
 
 ## Setup
 
+[Installation and demo video](https://youtu.be/uORnxt5DLTw)
+
 Download the webkeyboard.uf2 file from the latest
 [release](https://git.kkozai.com/kenji/webkeyboard/releases) and flash
 onto the Raspberry Pi Pico W by holding down the BOOTSEL button while plugging

websocket.c

@@ -13,6 +13,9 @@ static const char WS_RESPONSE[] = "HTTP/1.1 101 Switching Protocols\r\n" \
 	"Connection: Upgrade\r\n" \
 	"Sec-WebSocket-Accept: ";
 
+static uint8_t buf[WS_BUFFER_SIZE];
+static uint16_t buf_len=0;
+
 static tWSHandler ws_receive_cb = NULL;
 static tWSOpenHandler ws_open_cb = NULL;
 
@@ -330,30 +333,87 @@ static err_t ws_read(struct altcp_pcb *pcb, struct ws_state *wss, struct pbuf *p
 		// successful read, reset timeout
 		wss->retries = 0;
 
-		uint8_t mode = data[0] & 0x0F;
+		uint8_t fin = data[0] & 0x80;
+		uint8_t opcode = data[0] & 0x0F;
+		uint8_t masked = data[1] & 0x80;
 		uint16_t msg_len = data[1] & 0x7F;
-		switch (mode) {
+		uint8_t *msg;
-			case 0x01: //text
+		switch (msg_len) {
+			case 126: // next two bytes are length
+				memcpy(&msg_len, &data[2], 2);
+				if (len >= 8) {
+					msg = &data[8];
+				}
+				break;
+			case 127: // next four bytes are length
+				// lwIP's pbuf only handles 16-bit lengths, so error
+				return ERR_ARG;
+				
+				//memcpy(&msg_len, &data[2], 4);
+				//if (len >= 10) {
+				//	msg = &data[10];
+				//}
+				//break;
+			default:
+				if (len >= 6) {
+					msg = &data[6];
+				}
+				break;
+		}
+		switch (opcode) {
+			case OP_CONT:
+				LWIP_DEBUGF(WS_DEBUG, ("ws_read: received continuation frame\n"));
+			case OP_TEXT:
 				LWIP_DEBUGF(WS_DEBUG, ("ws_read: received text data\n"));
-			case 0x02: //binary
+			case OP_BINARY:
-				LWIP_DEBUGF(WS_DEBUG, ("ws_read: decoding data\n"));
+				LWIP_DEBUGF(WS_DEBUG, ("ws_read: decoding data, len=%u\n", msg_len));
-				if (len >= 6 && ws_receive_cb != NULL) {
+				if (msg && ws_receive_cb != NULL) {
+					// unmask the data if mask bit is received
+					if (masked) {
 						uint8_t *mask = &data[2];
-					uint8_t *msg = &data[6];
 	
 						for (int i=0; i<msg_len; i++) {
 							msg[i] ^= mask[i % 4];
 						}
+					} else {
+						// messages from client must be masked - disconnect
+						LWIP_DEBUGF(WS_DEBUG, ("ws_read: received unmasked message"));
+						return ERR_CLSD;
+					}
 					msg[msg_len]=0;
 
-					ws_receive_cb(msg, msg_len);
+					if (opcode != OP_CONT) { // not a continuation frame, reset buffer
+						buf_len=0;
+						memset(buf, 0x00, sizeof(buf));
+					}
+
+					memcpy(&buf[buf_len], msg, msg_len);
+					buf_len += msg_len;
+					
+					if (fin) { // last packet in message, process completed message
+						ws_receive_cb(buf, buf_len);
+					}
 				}
 				break;
-			case 0x08: //close
+			case OP_CLOSE:
 				LWIP_DEBUGF(WS_DEBUG, ("ws_read: close request"));
 				return ERR_CLSD;
+			case OP_PING:
+				// control frames cannot exceed 125 bytes in length
+				if (msg && msg_len <= 125) {
+					// send back a pong
+					uint8_t pong[2+msg_len];
+					pong[0]=0x8A;
+					pong[1]=msg_len;
+					memcpy(&pong[2], msg, msg_len);
+
+					return ws_send(wss, pong, msg_len+2);
+				}
+				return ERR_ARG;
+			case OP_PONG: // no response required for pong
+				return ERR_OK;
 			default:
-				LWIP_DEBUGF(WS_DEBUG, ("ws_read: invalid data mode %02X\n", mode));
+				LWIP_DEBUGF(WS_DEBUG, ("ws_read: invalid opcode %02X\n", opcode));
 				return ERR_ARG;
 		}
 

websocket.h

@@ -7,6 +7,14 @@
 #define WS_MAX_RETRIES		10
 #define WS_POLL_INTERVAL	60 // WS_POLL_INTERVAL/2 seconds
 #define WS_MAX_CONN			4
+#define WS_BUFFER_SIZE		512
+
+#define OP_CONT		0x00
+#define OP_TEXT		0x01
+#define OP_BINARY	0x02
+#define OP_CLOSE	0x08
+#define OP_PING		0x09
+#define OP_PONG		0x0A
 
 struct ws_state {
 	bool active;