From 59b9b3e26f20cd2e15f1b7cc1bc44d25b05affb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mariusz=20Kry=C5=84ski?= Date: Sat, 19 Nov 2016 16:35:20 +0100 Subject: [PATCH] support for custom websocket headers (#1573) Looks good to me. Thank you. Also: - allow for '\0's in received messages * add client:config for setting websocket headers Also: - headers are case-insensitive now * fix docs * fix typo * remove unnecessary luaL_argcheck calls * replace os_sprintf with simple string copy --- app/modules/websocket.c | 64 ++++++++++++++++++++--- app/websocket/websocketclient.c | 92 ++++++++++++++++++++++++++------- app/websocket/websocketclient.h | 8 ++- docs/en/modules/websocket.md | 25 +++++++-- 4 files changed, 160 insertions(+), 29 deletions(-) diff --git a/app/modules/websocket.c b/app/modules/websocket.c index 899d04dc..6feb6e82 100644 --- a/app/modules/websocket.c +++ b/app/modules/websocket.c @@ -15,6 +15,7 @@ #include "c_types.h" #include "c_string.h" +#include "c_stdlib.h" #include "websocketclient.h" @@ -45,7 +46,7 @@ static void websocketclient_onConnectionCallback(ws_info *ws) { } } -static void websocketclient_onReceiveCallback(ws_info *ws, char *message, int opCode) { +static void websocketclient_onReceiveCallback(ws_info *ws, int len, char *message, int opCode) { NODE_DBG("websocketclient_onReceiveCallback\n"); lua_State *L = lua_getstate(); @@ -59,7 +60,7 @@ static void websocketclient_onReceiveCallback(ws_info *ws, char *message, int op if (data->onReceive != LUA_NOREF) { lua_rawgeti(L, LUA_REGISTRYINDEX, data->onReceive); // load the callback function lua_rawgeti(L, LUA_REGISTRYINDEX, data->self_ref); // pass itself, #1 callback argument - lua_pushstring(L, message); // #2 callback argument + lua_pushlstring(L, message, len); // #2 callback argument lua_pushnumber(L, opCode); // #3 callback argument lua_call(L, 3, 0); } @@ -102,6 +103,7 @@ static int websocket_createClient(lua_State *L) { ws_info *ws = (ws_info *) lua_newuserdata(L, sizeof(ws_info)); ws->connectionState = 0; + ws->extraHeaders = NULL; ws->onConnection = &websocketclient_onConnectionCallback; ws->onReceive = &websocketclient_onReceiveCallback; ws->onFailure = &websocketclient_onCloseCallback; @@ -118,7 +120,6 @@ static int websocketclient_on(lua_State *L) { NODE_DBG("websocketclient_on\n"); ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT); - luaL_argcheck(L, ws, 1, "Client websocket expected"); ws_data *data = (ws_data *) ws->reservedData; @@ -170,7 +171,6 @@ static int websocketclient_connect(lua_State *L) { NODE_DBG("websocketclient_connect is called.\n"); ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT); - luaL_argcheck(L, ws, 1, "Client websocket expected"); ws_data *data = (ws_data *) ws->reservedData; @@ -188,11 +188,61 @@ static int websocketclient_connect(lua_State *L) { return 0; } +static header_t *realloc_headers(header_t *headers, int new_size) { + if(headers) { + for(header_t *header = headers; header->key; header++) { + c_free(header->value); + c_free(header->key); + } + c_free(headers); + } + if(!new_size) + return NULL; + return (header_t *)c_malloc(sizeof(header_t) * (new_size + 1)); +} + +static int websocketclient_config(lua_State *L) { + NODE_DBG("websocketclient_config is called.\n"); + + ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT); + + ws_data *data = (ws_data *) ws->reservedData; + + luaL_checktype(L, 2, LUA_TTABLE); + lua_getfield(L, 2, "headers"); + if(lua_istable(L, -1)) { + + lua_pushnil(L); + int size = 0; + while(lua_next(L, -2)) { + size++; + lua_pop(L, 1); + } + + ws->extraHeaders = realloc_headers(ws->extraHeaders, size); + if(ws->extraHeaders) { + header_t *header = ws->extraHeaders; + + lua_pushnil(L); + while(lua_next(L, -2)) { + header->key = c_strdup(lua_tostring(L, -2)); + header->value = c_strdup(lua_tostring(L, -1)); + header++; + lua_pop(L, 1); + } + + header->key = header->value = NULL; + } + } + lua_pop(L, 1); // pop headers + + return 0; +} + static int websocketclient_send(lua_State *L) { NODE_DBG("websocketclient_send is called.\n"); ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT); - luaL_argcheck(L, ws, 1, "Client websocket expected"); ws_data *data = (ws_data *) ws->reservedData; @@ -225,7 +275,8 @@ static int websocketclient_gc(lua_State *L) { NODE_DBG("websocketclient_gc\n"); ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT); - luaL_argcheck(L, ws, 1, "Client websocket expected"); + + ws->extraHeaders = realloc_headers(ws->extraHeaders, 0); ws_data *data = (ws_data *) ws->reservedData; @@ -265,6 +316,7 @@ static const LUA_REG_TYPE websocket_map[] = static const LUA_REG_TYPE websocketclient_map[] = { { LSTRKEY("on"), LFUNCVAL(websocketclient_on) }, + { LSTRKEY("config"), LFUNCVAL(websocketclient_config) }, { LSTRKEY("connect"), LFUNCVAL(websocketclient_connect) }, { LSTRKEY("send"), LFUNCVAL(websocketclient_send) }, { LSTRKEY("close"), LFUNCVAL(websocketclient_close) }, diff --git a/app/websocket/websocketclient.c b/app/websocket/websocketclient.c index d9fa29ce..834996e7 100644 --- a/app/websocket/websocketclient.c +++ b/app/websocket/websocketclient.c @@ -47,18 +47,10 @@ #define PORT_INSECURE 80 #define PORT_MAX_VALUE 65535 -// TODO: user agent configurable -#define WS_INIT_HEADERS "GET %s HTTP/1.1\r\n"\ - "Host: %s:%d\r\n"\ - "Upgrade: websocket\r\n"\ - "Connection: Upgrade\r\n"\ - "User-Agent: ESP8266\r\n"\ - "Sec-Websocket-Key: %s\r\n"\ - "Sec-WebSocket-Protocol: chat\r\n"\ - "Sec-WebSocket-Version: 13\r\n"\ - "\r\n" +#define WS_INIT_REQUEST "GET %s HTTP/1.1\r\n"\ + "Host: %s:%d\r\n" -#define WS_INIT_HEADERS_LENGTH 169 +#define WS_INIT_REQUEST_LENGTH 30 #define WS_GUID "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" #define WS_GUID_LENGTH 36 @@ -77,6 +69,13 @@ #define WS_OPCODE_PING 0x9 #define WS_OPCODE_PONG 0xA +header_t DEFAULT_HEADERS[] = { + {"User-Agent", "ESP8266"}, + {"Sec-WebSocket-Protocol", "chat"}, + {0} +}; +header_t *EMPTY_HEADERS = DEFAULT_HEADERS + sizeof(DEFAULT_HEADERS) / sizeof(header_t) - 1; + static char *cryptoSha1(char *data, unsigned int len) { SHA1_CTX ctx; SHA1Init(&ctx); @@ -128,6 +127,44 @@ static void generateSecKeys(char **key, char **expectedKey) { os_free(keyEncrypted); } +static char *_strcpy(char *dst, char *src) { + while(*dst++ = *src++); + return dst - 1; +} + +static int headers_length(header_t *headers) { + int length = 0; + for(; headers->key; headers++) + length += strlen(headers->key) + strlen(headers->value) + 4; + return length; +} + +static char *sprintf_headers(char *buf, ...) { + char *dst = buf; + va_list args; + va_start(args, buf); + for(header_t *header_set = va_arg(args, header_t *); header_set; header_set = va_arg(args, header_t *)) + for(header_t *header = header_set; header->key; header++) { + va_list args2; + va_start(args2, buf); + for(header_t *header_set2 = va_arg(args2, header_t *); header_set2; header_set2 = va_arg(args2, header_t *)) + for(header_t *header2 = header_set2; header2->key; header2++) { + if(header == header2) + goto ok; + if(!strcasecmp(header->key, header2->key)) + goto skip; + } +ok: + dst = _strcpy(dst, header->key); + dst = _strcpy(dst, ": "); + dst = _strcpy(dst, header->value); + dst = _strcpy(dst, "\r\n"); +skip:; + } + dst = _strcpy(dst, "\r\n"); + return dst; +} + static void ws_closeSentCallback(void *arg) { NODE_DBG("ws_closeSentCallback \n"); struct espconn *conn = (struct espconn *) arg; @@ -452,7 +489,7 @@ static void ws_receiveCallback(void *arg, char *buf, unsigned short len) { } else if (opCode == WS_OPCODE_PONG) { // ping alarm was already reset... } else { - if (ws->onReceive) ws->onReceive(ws, payload, opCode); + if (ws->onReceive) ws->onReceive(ws, payloadLength, payload, opCode); } os_free(payload); } @@ -509,7 +546,7 @@ static void ws_initReceiveCallback(void *arg, char *buf, unsigned short len) { } // Check server has valid sec key - if (strstr(buf, WS_HTTP_SEC_WEBSOCKET_ACCEPT) == NULL || strstr(buf, ws->expectedSecKey) == NULL) { + if (strstr(buf, ws->expectedSecKey) == NULL) { NODE_DBG("Server has invalid response\n"); ws->knownFailureCode = -7; if (ws->isSecure) @@ -550,12 +587,31 @@ static void connect_callback(void *arg) { char *key; generateSecKeys(&key, &ws->expectedSecKey); - char buf[WS_INIT_HEADERS_LENGTH + strlen(ws->path) + strlen(ws->hostname) + strlen(key)]; - int len = os_sprintf(buf, WS_INIT_HEADERS, ws->path, ws->hostname, ws->port, key); + header_t headers[] = { + {"Upgrade", "websocket"}, + {"Connection", "Upgrade"}, + {"Sec-WebSocket-Key", key}, + {"Sec-WebSocket-Version", "13"}, + {0} + }; + + header_t *extraHeaders = ws->extraHeaders ? ws->extraHeaders : EMPTY_HEADERS; + + char buf[WS_INIT_REQUEST_LENGTH + strlen(ws->path) + strlen(ws->hostname) + + headers_length(DEFAULT_HEADERS) + headers_length(headers) + headers_length(extraHeaders) + 2]; + + int len = os_sprintf( + buf, + WS_INIT_REQUEST, + ws->path, + ws->hostname, + ws->port + ); + + len = sprintf_headers(buf + len, headers, extraHeaders, DEFAULT_HEADERS, 0) - buf; os_free(key); - - NODE_DBG("connecting\n"); + NODE_DBG("request: %s", buf); if (ws->isSecure) espconn_secure_send(conn, (uint8_t *) buf, len); else @@ -630,7 +686,7 @@ static void dns_callback(const char *hostname, ip_addr_t *addr, void *arg) { ws_info *ws = (ws_info *) conn->reverse; if (ws->conn == NULL || ws->connectionState == 4) { - return; + return; } if (addr == NULL) { diff --git a/app/websocket/websocketclient.h b/app/websocket/websocketclient.h index 35746b48..0c029418 100644 --- a/app/websocket/websocketclient.h +++ b/app/websocket/websocketclient.h @@ -40,9 +40,14 @@ struct ws_info; typedef void (*ws_onConnectionCallback)(struct ws_info *wsInfo); -typedef void (*ws_onReceiveCallback)(struct ws_info *wsInfo, char *message, int opCode); +typedef void (*ws_onReceiveCallback)(struct ws_info *wsInfo, int len, char *message, int opCode); typedef void (*ws_onFailureCallback)(struct ws_info *wsInfo, int errorCode); +typedef struct { + char *key; + char *value; +} header_t; + typedef struct ws_info { int connectionState; @@ -51,6 +56,7 @@ typedef struct ws_info { int port; char *path; char *expectedSecKey; + header_t *extraHeaders; struct espconn *conn; void *reservedData; diff --git a/docs/en/modules/websocket.md b/docs/en/modules/websocket.md index 177c9abe..5f319c7c 100644 --- a/docs/en/modules/websocket.md +++ b/docs/en/modules/websocket.md @@ -7,10 +7,6 @@ A websocket *client* module that implements [RFC6455](https://tools.ietf.org/htm The implementation supports fragmented messages, automatically respondes to ping requests and periodically pings if the server isn't communicating. -!!! note - - Currently, it is **not** possible to change the request headers, most notably the user agent. - **SSL/TLS support** Take note of constraints documented in the [net module](net.md). @@ -70,6 +66,27 @@ ws = nil -- fully dispose the client as lua will now gc it ``` +## websocket.client:config(params) + +Configures websocket client instance. + +#### Syntax +`websocket:config(params)` + +#### Parameters +- `params` table with configuration parameters. Following keys are recognized: + - `headers` table of extra request headers affecting every request + +#### Returns +`nil` + +#### Example +```lua +ws = websocket.createClient() +ws:config({headers={['User-Agent']='NodeMCU'}}) +``` + + ## websocket.client:connect() Attempts to estabilish a websocket connection to the given URL.