websocket.c 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330
  1. // Module for websockets
  2. // Example usage:
  3. // ws = websocket.createClient()
  4. // ws:on("connection", function() ws:send('hi') end)
  5. // ws:on("receive", function(_, data, opcode) print(data) end)
  6. // ws:on("close", function(_, reasonCode) print('ws closed', reasonCode) end)
  7. // ws:connect('ws://echo.websocket.org')
  8. #include "lmem.h"
  9. #include "lauxlib.h"
  10. #include "platform.h"
  11. #include "module.h"
  12. #include <stdint.h>
  13. #include <string.h>
  14. #include <stddef.h>
  15. #include "websocket/websocketclient.h"
  16. #define METATABLE_WSCLIENT "websocket.client"
  17. typedef struct ws_data {
  18. int self_ref;
  19. int onConnection;
  20. int onReceive;
  21. int onClose;
  22. } ws_data;
  23. static void websocketclient_onConnectionCallback(ws_info *ws) {
  24. NODE_DBG("websocketclient_onConnectionCallback\n");
  25. lua_State *L = lua_getstate();
  26. if (ws == NULL || ws->reservedData == NULL) {
  27. luaL_error(L, "Client websocket is nil.\n");
  28. return;
  29. }
  30. ws_data *data = (ws_data *) ws->reservedData;
  31. if (data->onConnection != LUA_NOREF) {
  32. lua_rawgeti(L, LUA_REGISTRYINDEX, data->onConnection); // load the callback function
  33. lua_rawgeti(L, LUA_REGISTRYINDEX, data->self_ref); // pass itself, #1 callback argument
  34. luaL_pcallx(L, 1, 0);
  35. }
  36. }
  37. static void websocketclient_onReceiveCallback(ws_info *ws, int len, char *message, int opCode) {
  38. NODE_DBG("websocketclient_onReceiveCallback\n");
  39. lua_State *L = lua_getstate();
  40. if (ws == NULL || ws->reservedData == NULL) {
  41. luaL_error(L, "Client websocket is nil.\n");
  42. return;
  43. }
  44. ws_data *data = (ws_data *) ws->reservedData;
  45. if (data->onReceive != LUA_NOREF) {
  46. lua_rawgeti(L, LUA_REGISTRYINDEX, data->onReceive); // load the callback function
  47. lua_rawgeti(L, LUA_REGISTRYINDEX, data->self_ref); // pass itself, #1 callback argument
  48. lua_pushlstring(L, message, len); // #2 callback argument
  49. lua_pushnumber(L, opCode); // #3 callback argument
  50. luaL_pcallx(L, 3, 0);
  51. }
  52. }
  53. static void websocketclient_onCloseCallback(ws_info *ws, int errorCode) {
  54. NODE_DBG("websocketclient_onCloseCallback\n");
  55. lua_State *L = lua_getstate();
  56. if (ws == NULL || ws->reservedData == NULL) {
  57. luaL_error(L, "Client websocket is nil.\n");
  58. return;
  59. }
  60. ws_data *data = (ws_data *) ws->reservedData;
  61. if (data->onClose != LUA_NOREF) {
  62. lua_rawgeti(L, LUA_REGISTRYINDEX, data->onClose); // load the callback function
  63. lua_rawgeti(L, LUA_REGISTRYINDEX, data->self_ref); // pass itself, #1 callback argument
  64. lua_pushnumber(L, errorCode); // pass the error code, #2 callback argument
  65. luaL_pcallx(L, 2, 0);
  66. }
  67. // free self-reference to allow gc (no futher callback will be called until next ws:connect())
  68. lua_gc(L, LUA_GCSTOP, 0); // required to avoid freeing ws_data
  69. luaL_unref(L, LUA_REGISTRYINDEX, data->self_ref);
  70. data->self_ref = LUA_NOREF;
  71. lua_gc(L, LUA_GCRESTART, 0);
  72. }
  73. static int websocket_createClient(lua_State *L) {
  74. NODE_DBG("websocket_createClient\n");
  75. // create user data
  76. ws_data *data = (ws_data *) luaM_malloc(L, sizeof(ws_data));
  77. data->onConnection = LUA_NOREF;
  78. data->onReceive = LUA_NOREF;
  79. data->onClose = LUA_NOREF;
  80. data->self_ref = LUA_NOREF; // only set when ws:connect is called
  81. ws_info *ws = (ws_info *) lua_newuserdata(L, sizeof(ws_info));
  82. ws->connectionState = 0;
  83. ws->extraHeaders = NULL;
  84. ws->onConnection = &websocketclient_onConnectionCallback;
  85. ws->onReceive = &websocketclient_onReceiveCallback;
  86. ws->onFailure = &websocketclient_onCloseCallback;
  87. ws->reservedData = data;
  88. // set its metatable
  89. luaL_getmetatable(L, METATABLE_WSCLIENT);
  90. lua_setmetatable(L, -2);
  91. return 1;
  92. }
  93. static int websocketclient_on(lua_State *L) {
  94. NODE_DBG("websocketclient_on\n");
  95. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  96. ws_data *data = (ws_data *) ws->reservedData;
  97. int handle = luaL_checkoption(L, 2, NULL, (const char * const[]){ "connection", "receive", "close", NULL });
  98. luaL_argcheck(L, lua_isnil(L,3) || lua_isfunction(L, 3), 3, "function or nil");
  99. switch (handle) {
  100. case 0:
  101. NODE_DBG("connection\n");
  102. luaL_unref(L, LUA_REGISTRYINDEX, data->onConnection);
  103. data->onConnection = LUA_NOREF;
  104. if (!lua_isnil(L,3)) {
  105. lua_pushvalue(L, 3); // copy argument (func) to the top of stack
  106. data->onConnection = luaL_ref(L, LUA_REGISTRYINDEX);
  107. }
  108. break;
  109. case 1:
  110. NODE_DBG("receive\n");
  111. luaL_unref(L, LUA_REGISTRYINDEX, data->onReceive);
  112. data->onReceive = LUA_NOREF;
  113. if (!lua_isnil(L,3)) {
  114. lua_pushvalue(L, 3); // copy argument (func) to the top of stack
  115. data->onReceive = luaL_ref(L, LUA_REGISTRYINDEX);
  116. }
  117. break;
  118. case 2:
  119. NODE_DBG("close\n");
  120. luaL_unref(L, LUA_REGISTRYINDEX, data->onClose);
  121. data->onClose = LUA_NOREF;
  122. if (!lua_isnil(L,3)) {
  123. lua_pushvalue(L, 3); // copy argument (func) to the top of stack
  124. data->onClose = luaL_ref(L, LUA_REGISTRYINDEX);
  125. }
  126. break;
  127. }
  128. return 0;
  129. }
  130. static int websocketclient_connect(lua_State *L) {
  131. NODE_DBG("websocketclient_connect is called.\n");
  132. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  133. ws_data *data = (ws_data *) ws->reservedData;
  134. if (ws->connectionState != 0 && ws->connectionState != 4) {
  135. return luaL_error(L, "Websocket already connecting or connected.\n");
  136. }
  137. ws->connectionState = 0;
  138. lua_pushvalue(L, 1); // copy userdata to the top of stack to allow ref
  139. data->self_ref = luaL_ref(L, LUA_REGISTRYINDEX);
  140. const char *url = luaL_checkstring(L, 2);
  141. ws_connect(ws, url);
  142. return 0;
  143. }
  144. static header_t *realloc_headers(header_t *headers, int new_size) {
  145. if(headers) {
  146. for(header_t *header = headers; header->key; header++) {
  147. free(header->value);
  148. free(header->key);
  149. }
  150. free(headers);
  151. }
  152. if(!new_size)
  153. return NULL;
  154. return (header_t *)malloc(sizeof(header_t) * (new_size + 1));
  155. }
  156. static int websocketclient_config(lua_State *L) {
  157. NODE_DBG("websocketclient_config is called.\n");
  158. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  159. ws_data *data = (ws_data *) ws->reservedData;
  160. luaL_checktype(L, 2, LUA_TTABLE);
  161. lua_getfield(L, 2, "headers");
  162. if(lua_istable(L, -1)) {
  163. lua_pushnil(L);
  164. int size = 0;
  165. while(lua_next(L, -2)) {
  166. size++;
  167. lua_pop(L, 1);
  168. }
  169. ws->extraHeaders = realloc_headers(ws->extraHeaders, size);
  170. if(ws->extraHeaders) {
  171. header_t *header = ws->extraHeaders;
  172. lua_pushnil(L);
  173. while(lua_next(L, -2)) {
  174. header->key = strdup(lua_tostring(L, -2));
  175. header->value = strdup(lua_tostring(L, -1));
  176. header++;
  177. lua_pop(L, 1);
  178. }
  179. header->key = header->value = NULL;
  180. }
  181. }
  182. lua_pop(L, 1); // pop headers
  183. return 0;
  184. }
  185. static int websocketclient_send(lua_State *L) {
  186. NODE_DBG("websocketclient_send is called.\n");
  187. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  188. ws_data *data = (ws_data *) ws->reservedData;
  189. if (ws->connectionState != 3) {
  190. // should this be an onFailure callback instead?
  191. return luaL_error(L, "Websocket isn't connected.\n");
  192. }
  193. int msgLength;
  194. const char *msg = luaL_checklstring(L, 2, &msgLength);
  195. int opCode = 1; // default: text message
  196. if (lua_gettop(L) == 3) {
  197. opCode = luaL_checkint(L, 3);
  198. }
  199. ws_send(ws, opCode, msg, (unsigned short) msgLength);
  200. return 0;
  201. }
  202. static int websocketclient_close(lua_State *L) {
  203. NODE_DBG("websocketclient_close.\n");
  204. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  205. ws_close(ws);
  206. return 0;
  207. }
  208. static int websocketclient_gc(lua_State *L) {
  209. NODE_DBG("websocketclient_gc\n");
  210. ws_info *ws = (ws_info *) luaL_checkudata(L, 1, METATABLE_WSCLIENT);
  211. ws->extraHeaders = realloc_headers(ws->extraHeaders, 0);
  212. ws_data *data = (ws_data *) ws->reservedData;
  213. luaL_unref(L, LUA_REGISTRYINDEX, data->onConnection);
  214. luaL_unref(L, LUA_REGISTRYINDEX, data->onReceive);
  215. if (data->onClose != LUA_NOREF) {
  216. if (ws->connectionState != 4) { // only call if connection open
  217. lua_rawgeti(L, LUA_REGISTRYINDEX, data->onClose);
  218. lua_pushnumber(L, -100);
  219. luaL_pcallx(L, 1, 0);
  220. }
  221. luaL_unref(L, LUA_REGISTRYINDEX, data->onClose);
  222. }
  223. if (data->self_ref != LUA_NOREF) {
  224. lua_gc(L, LUA_GCSTOP, 0); // required to avoid freeing ws_data
  225. luaL_unref(L, LUA_REGISTRYINDEX, data->self_ref);
  226. data->self_ref = LUA_NOREF;
  227. lua_gc(L, LUA_GCRESTART, 0);
  228. }
  229. NODE_DBG("freeing lua data\n");
  230. luaM_free(L, data);
  231. NODE_DBG("done freeing lua data\n");
  232. return 0;
  233. }
  234. LROT_BEGIN(websocket, NULL, 0)
  235. LROT_FUNCENTRY( createClient, websocket_createClient )
  236. LROT_END(websocket, NULL, 0)
  237. LROT_BEGIN(websocketclient, NULL, LROT_MASK_GC_INDEX)
  238. LROT_FUNCENTRY( __gc, websocketclient_gc )
  239. LROT_TABENTRY( __index, websocketclient )
  240. LROT_FUNCENTRY( on, websocketclient_on )
  241. LROT_FUNCENTRY( config, websocketclient_config )
  242. LROT_FUNCENTRY( connect, websocketclient_connect )
  243. LROT_FUNCENTRY( send, websocketclient_send )
  244. LROT_FUNCENTRY( close, websocketclient_close )
  245. LROT_END(websocketclient, NULL, LROT_MASK_GC_INDEX)
  246. int loadWebsocketModule(lua_State *L) {
  247. luaL_rometatable(L, METATABLE_WSCLIENT, LROT_TABLEREF( websocketclient));
  248. return 0;
  249. }
  250. NODEMCU_MODULE(WEBSOCKET, "websocket", websocket, loadWebsocketModule);