initial commit
This commit is contained in:
+470
@@ -0,0 +1,470 @@
|
||||
#include <string.h>
|
||||
|
||||
#include "lwip/altcp.h"
|
||||
#include "lwip/debug.h"
|
||||
#include "mbedtls/base64.h"
|
||||
#include "mbedtls/sha1.h"
|
||||
|
||||
#include "websocket.h"
|
||||
|
||||
static const char WS_GUID[] = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
|
||||
static const char WS_RESPONSE[] = "HTTP/1.1 101 Switching Protocols\r\n" \
|
||||
"Upgrade: websocket\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;
|
||||
|
||||
static struct ws_state * ws_connections;
|
||||
static uint8_t ws_num_conns = 0;
|
||||
|
||||
static struct ws_state* ws_state_alloc(void);
|
||||
static void ws_state_init(struct ws_state *wss);
|
||||
static void ws_state_free(struct ws_state *wss);
|
||||
static void ws_server_init_pcb( struct altcp_pcb *pcb, uint16_t port);
|
||||
static err_t ws_accept(void *arg, struct altcp_pcb *pcb, err_t err);
|
||||
static err_t ws_recv(void *arg, struct altcp_pcb *pcb, struct pbuf *p, err_t err);
|
||||
static err_t ws_sent(void *arg, struct altcp_pcb *pcb, uint16_t len);
|
||||
static void ws_err (void *arg, err_t err);
|
||||
static err_t ws_close_conn(struct altcp_pcb *pcb, struct ws_state *wss);
|
||||
static err_t ws_close_or_abort_conn(struct altcp_pcb *pcb, struct ws_state *wss, uint8_t abort_conn);
|
||||
static err_t ws_poll(void *arg, struct altcp_pcb *pcb);
|
||||
static err_t ws_handshake(struct altcp_pcb *pcb, struct ws_state *wss, struct pbuf *p);
|
||||
static err_t ws_read(struct altcp_pcb *pcb, struct ws_state *wss, struct pbuf *p);
|
||||
static err_t ws_send(struct ws_state *wss, uint8_t *data, uint16_t len);
|
||||
|
||||
// allocate memory for ws_state instance
|
||||
static struct ws_state * ws_state_alloc(void) {
|
||||
struct ws_state *ret = WS_ALLOC_WS_STATE();
|
||||
|
||||
if ( ret != NULL) {
|
||||
ws_state_init(ret);
|
||||
if (ws_connections == NULL) {
|
||||
ws_connections = ret;
|
||||
} else {
|
||||
struct ws_state *last;
|
||||
for (last=ws_connections; last->next != NULL; last=last->next);
|
||||
LWIP_ASSERT("last != NULL", last != NULL);
|
||||
last->next = ret;
|
||||
}
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
// initiate ws_state instance
|
||||
static void ws_state_init(struct ws_state *wss) {
|
||||
memset(wss, 0, sizeof(struct ws_state));
|
||||
wss->active = false;
|
||||
}
|
||||
|
||||
// free memory from ws_state instance
|
||||
static void ws_state_free(struct ws_state *wss) {
|
||||
if (wss != NULL) {
|
||||
if (ws_connections == wss) {
|
||||
ws_connections = wss->next;
|
||||
} else {
|
||||
struct ws_state * last;
|
||||
for (last = ws_connections; last->next != NULL; last = last->next) {
|
||||
if (last->next == wss) {
|
||||
last->next = wss->next;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
mem_free(wss);
|
||||
}
|
||||
}
|
||||
|
||||
// initiate websocket server on specified pcb
|
||||
static void ws_server_init_pcb( struct altcp_pcb *pcb, uint16_t port) {
|
||||
err_t err;
|
||||
|
||||
if (pcb) {
|
||||
altcp_setprio(pcb, TCP_PRIO_MIN);
|
||||
err = altcp_bind(pcb, IP_ANY_TYPE, port);
|
||||
LWIP_UNUSED_ARG(err);
|
||||
LWIP_ASSERT("ws_server_init: tcp_bind failed", err == ERR_OK);
|
||||
pcb = altcp_listen(pcb);
|
||||
LWIP_ASSERT("ws_server_init: tcp_listen failed", pcb != NULL);
|
||||
altcp_accept(pcb, ws_accept);
|
||||
}
|
||||
}
|
||||
|
||||
// initiate a websocket server
|
||||
void ws_server_init(void) {
|
||||
struct altcp_pcb *pcb = altcp_tcp_new_ip_type(IPADDR_TYPE_ANY);
|
||||
LWIP_ASSERT("ws_server_init: tcp_new failed", pcb != NULL);
|
||||
ws_server_init_pcb(pcb, WS_PORT);
|
||||
}
|
||||
|
||||
// set ws_receive_handler
|
||||
void ws_set_receive_handler( tWSHandler ws_handler)
|
||||
{
|
||||
ws_receive_cb = ws_handler;
|
||||
}
|
||||
|
||||
// set ws_open_handler
|
||||
void ws_set_open_handler( tWSOpenHandler ws_handler)
|
||||
{
|
||||
ws_open_cb = ws_handler;
|
||||
}
|
||||
|
||||
// callback for accepted websocket connection
|
||||
static err_t ws_accept(void *arg, struct altcp_pcb *pcb, err_t err) {
|
||||
struct ws_state *wss;
|
||||
LWIP_UNUSED_ARG(err);
|
||||
LWIP_UNUSED_ARG(arg);
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_accept %p / %p\n", (void *)pcb, arg));
|
||||
|
||||
if ((err != ERR_OK) || (pcb == NULL)) {
|
||||
return ERR_VAL;
|
||||
}
|
||||
|
||||
// create new ws_state object
|
||||
wss = ws_state_alloc();
|
||||
if (wss == NULL) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_accept: Out of memory, RST\n"));
|
||||
return ERR_MEM;
|
||||
}
|
||||
wss->pcb = pcb;
|
||||
|
||||
// make ws_state object the argument of callbacks
|
||||
altcp_arg(pcb, wss);
|
||||
|
||||
// register callbacks for tcp events
|
||||
altcp_recv(pcb, ws_recv);
|
||||
altcp_sent(pcb, ws_sent);
|
||||
altcp_poll(pcb, ws_poll, WS_POLL_INTERVAL);
|
||||
altcp_err(pcb, ws_err);
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
// call when data is received
|
||||
static err_t ws_recv(void *arg, struct altcp_pcb *pcb, struct pbuf *p, err_t err) {
|
||||
struct ws_state *wss = (struct ws_state *) arg;
|
||||
|
||||
if ((err != ERR_OK) || (p == NULL) || (wss == NULL)) {
|
||||
// error or closed by client
|
||||
if (p != NULL) {
|
||||
// inform TCP that we have taken the data
|
||||
altcp_recved(pcb, p->tot_len);
|
||||
pbuf_free(p);
|
||||
}
|
||||
if (wss == NULL) {
|
||||
// should not occur
|
||||
LWIP_DEBUGF(WS_DEBUG, ("Error, ws_recv: wss is NULL, close\n"));
|
||||
}
|
||||
ws_close_conn(pcb, wss);
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
if (wss->active) {
|
||||
// process websocket message
|
||||
err = ws_read(pcb, wss, p);
|
||||
} else {
|
||||
// init websocket connection
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_recv: websocket inactive, checking for handshake\n"));
|
||||
|
||||
err = ws_handshake(pcb, wss, p);
|
||||
}
|
||||
|
||||
// inform TCP that we have taken the data.
|
||||
altcp_recved(pcb, p->tot_len);
|
||||
pbuf_free(p);
|
||||
|
||||
if (err == ERR_CLSD) {
|
||||
ws_close_conn(pcb, wss);
|
||||
}
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
// called when data has been sent over the websocket
|
||||
static err_t ws_sent(void *arg, struct altcp_pcb *pcb, uint16_t len) {
|
||||
(void) pcb;
|
||||
|
||||
struct ws_state *wss = (struct ws_state *)arg;
|
||||
|
||||
LWIP_DEBUGF(WS_DEBUG | LWIP_DBG_TRACE, ("ws_sent %p\n", (void*) pcb));
|
||||
|
||||
LWIP_UNUSED_ARG(len);
|
||||
|
||||
if (wss == NULL) {
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
wss->retries = 0;
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
// called when there is a websocket error
|
||||
static void ws_err (void *arg, err_t err) {
|
||||
struct ws_state *wss = (struct ws_state *) arg;
|
||||
LWIP_UNUSED_ARG(err);
|
||||
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_err: %s", lwip_strerr(err)));
|
||||
|
||||
if (wss != NULL) {
|
||||
ws_state_free(wss);
|
||||
}
|
||||
}
|
||||
|
||||
// initiate close of connection
|
||||
static err_t ws_close_conn(struct altcp_pcb *pcb, struct ws_state *wss) {
|
||||
return ws_close_or_abort_conn(pcb, wss, 0);
|
||||
}
|
||||
|
||||
// call when closing connection or connection was aborted
|
||||
static err_t ws_close_or_abort_conn(struct altcp_pcb *pcb, struct ws_state *wss,
|
||||
uint8_t abort_conn) {
|
||||
|
||||
err_t err;
|
||||
LWIP_DEBUGF(WS_DEBUG, ("Closing connection %p\n", (void *)pcb));
|
||||
|
||||
// clear callbacks
|
||||
altcp_arg(pcb, NULL);
|
||||
altcp_recv(pcb, NULL);
|
||||
altcp_sent(pcb, NULL);
|
||||
altcp_poll(pcb, NULL, 0);
|
||||
altcp_err(pcb, NULL);
|
||||
|
||||
// remove and free memory from ws_state object
|
||||
if (wss != NULL) {
|
||||
ws_state_free(wss);
|
||||
}
|
||||
|
||||
if (abort_conn) {
|
||||
altcp_abort(pcb);
|
||||
return ERR_OK;
|
||||
}
|
||||
err = altcp_close(pcb);
|
||||
if (err != ERR_OK) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("Error %d closing %p\n", err, (void *)pcb));
|
||||
// error closing, try again later in poll
|
||||
altcp_poll(pcb, ws_poll, WS_POLL_INTERVAL);
|
||||
}
|
||||
return err;
|
||||
}
|
||||
|
||||
// callback for polling process
|
||||
static err_t ws_poll(void *arg, struct altcp_pcb *pcb) {
|
||||
struct ws_state *wss = (struct ws_state *) arg;
|
||||
if (wss == NULL) {
|
||||
err_t closed;
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_poll: arg is NULL, close\n"));
|
||||
closed = ws_close_conn(pcb, NULL);
|
||||
LWIP_UNUSED_ARG(closed);
|
||||
if (closed == ERR_MEM) {
|
||||
altcp_abort(pcb);
|
||||
return ERR_ABRT;
|
||||
}
|
||||
return ERR_OK;
|
||||
} else {
|
||||
wss->retries++;
|
||||
if (wss->retries == WS_MAX_RETRIES) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_poll: too may retries, close\n"));
|
||||
ws_close_conn(pcb, wss);
|
||||
return ERR_OK;
|
||||
}
|
||||
}
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
|
||||
// check for and complete handshake with client
|
||||
static err_t ws_handshake(struct altcp_pcb *pcb, struct ws_state *wss, struct pbuf *p){
|
||||
uint8_t *data = (uint8_t *) p->payload;
|
||||
uint16_t len = p->len;
|
||||
|
||||
// check if client is initiating a websocket connecttion
|
||||
if (strstr(data, "Upgrade: websocket")) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_handshake: received websocket upgrade request\n"));
|
||||
|
||||
// search for websocket security key
|
||||
char *key_start = strstr(data, "Sec-WebSocket-Key: ");
|
||||
|
||||
if (key_start) {
|
||||
key_start += 19;
|
||||
const char *key_end = strstr(key_start, "\r\n");
|
||||
if (key_end) {
|
||||
char key[64];
|
||||
uint16_t key_len = key_end-key_start;
|
||||
if ( (key_len>0) && (key_len + sizeof(WS_GUID) < sizeof(key)) ) {
|
||||
// create response key by concatenating with websocket GUID,
|
||||
// taking SHA1 hash, then encoding in base 64
|
||||
strncpy(key, key_start, key_len);
|
||||
strlcpy(&key[key_len], WS_GUID, sizeof(key)-key_len);
|
||||
|
||||
key_len += sizeof(WS_GUID)-1;
|
||||
unsigned char key_sha1[20];
|
||||
unsigned char key_base64[29];
|
||||
size_t encoded_len;
|
||||
mbedtls_sha1( (unsigned char *) key, key_len, key_sha1);
|
||||
mbedtls_base64_encode( key_base64, 29, &encoded_len, key_sha1, 20);
|
||||
|
||||
// create response packet with encoded response key
|
||||
unsigned char response[sizeof(WS_RESPONSE) + sizeof(key_base64)+3];
|
||||
size_t count = sprintf(response, "%s%s\r\n\r\n", WS_RESPONSE, key_base64);
|
||||
|
||||
// send completed data packet
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_handshake: sending response\n"));
|
||||
if(altcp_write(pcb, response, count, TCP_WRITE_FLAG_COPY) == ERR_OK) {
|
||||
wss->active = true;
|
||||
}
|
||||
|
||||
if (ws_open_cb != NULL) {
|
||||
ws_open_cb(wss);
|
||||
}
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
}
|
||||
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_handshake: key overflow\n"));
|
||||
return ERR_MEM;
|
||||
} else {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_handshake: key not received\n"));
|
||||
return ERR_ARG;
|
||||
}
|
||||
}
|
||||
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_handshake: not a websocket request\n"));
|
||||
return ERR_ARG;
|
||||
}
|
||||
|
||||
// handle reading of websocket data and pass to ws_receive_cb
|
||||
static err_t ws_read(struct altcp_pcb *pcb, struct ws_state *wss, struct pbuf *p) {
|
||||
(void) pcb;
|
||||
uint8_t *data = (uint8_t *) p->payload;
|
||||
uint16_t len = p->len;
|
||||
|
||||
if (data != NULL && len > 1) {
|
||||
// successful read, reset timeout
|
||||
wss->retries = 0;
|
||||
|
||||
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;
|
||||
uint8_t *msg;
|
||||
|
||||
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
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_read: received 64-bit length %u\n", msg_len));
|
||||
return ERR_MEM;
|
||||
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 OP_BINARY:
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_read: decoding data, len=%u\n", msg_len));
|
||||
if (msg && ws_receive_cb != NULL) {
|
||||
// unmask the data if mask bit is received
|
||||
if (masked) {
|
||||
uint8_t *mask = &data[2];
|
||||
|
||||
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;
|
||||
|
||||
if (opcode != OP_CONT) { // not a continuation frame, reset buffer
|
||||
buf_len=0;
|
||||
memset(buf, 0x00, sizeof(buf));
|
||||
}
|
||||
|
||||
if (buf_len + msg_len > WS_BUFFER_SIZE) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_read: message exceeds buffer size %u+%u\n", buf_len, msg_len));
|
||||
return ERR_MEM;
|
||||
}
|
||||
|
||||
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 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 opcode %02X\n", opcode));
|
||||
return ERR_ARG;
|
||||
}
|
||||
|
||||
return ERR_OK;
|
||||
}
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_read: received empty payload\n"));
|
||||
return ERR_VAL;
|
||||
}
|
||||
|
||||
static err_t ws_send(struct ws_state *wss, uint8_t *data, uint16_t len) {
|
||||
uint8_t buf[128];
|
||||
buf[0] = 0x81;
|
||||
buf[1] = len & 0x7F;
|
||||
memcpy(&buf[2], data, len);
|
||||
|
||||
err_t err;
|
||||
err = altcp_write(wss->pcb, buf, len+2, TCP_WRITE_FLAG_COPY);
|
||||
if (err == ERR_OK) {
|
||||
altcp_output(wss->pcb);
|
||||
}
|
||||
|
||||
return err;
|
||||
}
|
||||
|
||||
void ws_send_all(uint8_t *data, uint16_t len) {
|
||||
// send message to all connections
|
||||
if (ws_connections != NULL) {
|
||||
struct ws_state *wss;
|
||||
err_t err;
|
||||
for (wss=ws_connections; wss != NULL; wss=wss->next) {
|
||||
err = ws_send(wss, data, len);
|
||||
if (err != ERR_OK ) {
|
||||
LWIP_DEBUGF(WS_DEBUG, ("ws_send_all: error sending to %p\n", wss));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user