websocket.c 9.4 KB

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