1// Aseprite
2// Copyright (C) 2021 Igara Studio S.A.
3//
4// This program is distributed under the terms of
5// the End-User License Agreement for Aseprite.
6
7#ifdef HAVE_CONFIG_H
8#include "config.h"
9#endif
10
11#include "app/app.h"
12#include "app/console.h"
13#include "app/script/engine.h"
14#include "app/script/luacpp.h"
15#include "app/script/security.h"
16#include "ui/timer.h"
17#include "ui/manager.h"
18#include "ui/system.h"
19
20#include <ixwebsocket/IXNetSystem.h>
21#include <ixwebsocket/IXWebSocket.h>
22#include <sstream>
23#include <set>
24
25namespace app {
26namespace script {
27
28namespace {
29
30// Additional "enum" value to make message callback simpler
31#define MESSAGE_TYPE_BINARY ((int)ix::WebSocketMessageType::Fragment + 10)
32
33static std::unique_ptr<ui::Timer> g_timer;
34static std::set<ix::WebSocket*> g_connections;
35
36static void close_ws(ix::WebSocket* ws)
37{
38 ws->stop();
39
40 g_connections.erase(ws);
41 if (g_connections.empty())
42 g_timer.reset();
43}
44
45int WebSocket_new(lua_State* L)
46{
47 static std::once_flag f;
48 std::call_once(f, &ix::initNetSystem);
49
50 auto ws = new ix::WebSocket();
51
52 push_ptr(L, ws);
53
54 if (lua_istable(L, 1)) {
55 lua_getfield(L, 1, "url");
56 if (const char* s = lua_tostring(L, -1)) {
57 if (!ask_access(L, s, FileAccessMode::OpenSocket, ResourceType::WebSocket))
58 return luaL_error(L, "the script doesn't have access to create a WebSocket for '%s'", s);
59
60 ws->setUrl(s);
61 }
62 lua_pop(L, 1);
63
64 lua_getfield(L, 1, "deflate");
65 if (lua_toboolean(L, -1)) {
66 ws->enablePerMessageDeflate();
67 }
68 else {
69 ws->disablePerMessageDeflate();
70 }
71 lua_pop(L, 1);
72
73 int type = lua_getfield(L, 1, "minreconnectwait");
74 if (type == LUA_TNUMBER) {
75 ws->setMinWaitBetweenReconnectionRetries(1000 * lua_tonumber(L, -1));
76 }
77 lua_pop(L, 1);
78
79 type = lua_getfield(L, 1, "maxreconnectwait");
80 if (type == LUA_TNUMBER) {
81 ws->setMaxWaitBetweenReconnectionRetries(1000 * lua_tonumber(L, -1));
82 }
83 lua_pop(L, 1);
84
85 type = lua_getfield(L, 1, "onreceive");
86 if (type == LUA_TFUNCTION) {
87 int onreceiveRef = luaL_ref(L, LUA_REGISTRYINDEX);
88
89 ws->setOnMessageCallback(
90 [L, ws, onreceiveRef](const ix::WebSocketMessagePtr& msg) {
91 int msgType =
92 (msg->binary ? MESSAGE_TYPE_BINARY : static_cast<int>(msg->type));
93 std::string msgData = msg->str;
94
95 ui::execute_from_ui_thread([=]() {
96 lua_rawgeti(L, LUA_REGISTRYINDEX, onreceiveRef);
97 lua_pushinteger(L, msgType);
98 lua_pushlstring(L, msgData.c_str(), msgData.length());
99
100 if (lua_pcall(L, 2, 0, 0)) {
101 if (const char* s = lua_tostring(L, -1)) {
102 App::instance()->scriptEngine()->consolePrint(s);
103 ws->stop();
104 }
105 }
106 });
107 });
108 }
109 else {
110 // Set a default handler to avoid a std::bad_function_call exception
111 ws->setOnMessageCallback([](const ix::WebSocketMessagePtr& msg) { });
112 lua_pop(L, 1);
113 }
114 }
115
116 return 1;
117}
118
119int WebSocket_gc(lua_State* L)
120{
121 auto ws = get_ptr<ix::WebSocket>(L, 1);
122 close_ws(ws);
123 delete ws;
124 return 0;
125}
126
127int WebSocket_sendText(lua_State* L)
128{
129 auto ws = get_ptr<ix::WebSocket>(L, 1);
130
131 if (ws->getReadyState() != ix::ReadyState::Open) {
132 return luaL_error(L, "WebSocket is not connected, can't send text");
133 }
134
135 std::stringstream data;
136 int argc = lua_gettop(L);
137
138 for (int i = 2; i <= argc; i++) {
139 size_t bufLen;
140 const char* buf = lua_tolstring(L, i, &bufLen);
141 data.write(buf, bufLen);
142 }
143
144 if (!ws->sendText(data.str()).success) {
145 return luaL_error(L, "WebSocket failed to send text");
146 }
147 return 0;
148}
149
150int WebSocket_sendBinary(lua_State* L)
151{
152 auto ws = get_ptr<ix::WebSocket>(L, 1);
153
154 if (ws->getReadyState() != ix::ReadyState::Open) {
155 return luaL_error(L, "WebSocket is not connected, can't send data");
156 }
157
158 std::stringstream data;
159 int argc = lua_gettop(L);
160
161 for (int i = 2; i <= argc; i++) {
162 size_t bufLen;
163 const char* buf = lua_tolstring(L, i, &bufLen);
164 data.write(buf, bufLen);
165 }
166
167 if (!ws->sendBinary(data.str()).success) {
168 return luaL_error(L, "WebSocket failed to send data");
169 }
170 return 0;
171}
172
173int WebSocket_sendPing(lua_State* L)
174{
175 auto ws = get_ptr<ix::WebSocket>(L, 1);
176
177 if (ws->getReadyState() != ix::ReadyState::Open) {
178 return luaL_error(L, "WebSocket is not connected, can't send ping");
179 }
180
181 size_t bufLen;
182 const char* buf = lua_tolstring(L, 2, &bufLen);
183 std::string data(buf, bufLen);
184
185 if (!ws->ping(data).success) {
186 return luaL_error(L, "WebSocket failed to send ping");
187 }
188 return 0;
189}
190
191int WebSocket_connect(lua_State* L)
192{
193 auto ws = get_ptr<ix::WebSocket>(L, 1);
194 ws->start();
195
196 if (g_connections.empty()) {
197#ifdef ENABLE_UI
198 if (App::instance()->isGui()) {
199 g_timer = std::make_unique<ui::Timer>(33, ui::Manager::getDefault());
200 g_timer->start();
201 }
202#endif
203 }
204 g_connections.insert(ws);
205
206 return 0;
207}
208
209int WebSocket_close(lua_State* L)
210{
211 auto ws = get_ptr<ix::WebSocket>(L, 1);
212 close_ws(ws);
213 return 0;
214}
215
216int WebSocket_get_url(lua_State* L)
217{
218 auto ws = get_ptr<ix::WebSocket>(L, 1);
219 lua_pushstring(L, ws->getUrl().c_str());
220 return 1;
221}
222
223const luaL_Reg WebSocket_methods[] = {
224 { "__gc", WebSocket_gc },
225 { "close", WebSocket_close },
226 { "connect", WebSocket_connect },
227 { "sendText", WebSocket_sendText },
228 { "sendBinary", WebSocket_sendBinary },
229 { "sendPing", WebSocket_sendPing },
230 { nullptr, nullptr }
231};
232
233const Property WebSocket_properties[] = {
234 { "url", WebSocket_get_url, nullptr },
235 { nullptr, nullptr, nullptr }
236};
237
238} // anonymous namespace
239
240using WebSocket = ix::WebSocket;
241DEF_MTNAME(WebSocket);
242
243void register_websocket_class(lua_State* L)
244{
245 REG_CLASS(L, WebSocket);
246 REG_CLASS_NEW(L, WebSocket);
247 REG_CLASS_PROPERTIES(L, WebSocket);
248
249 // message type enum
250 lua_newtable(L);
251 lua_pushvalue(L, -1);
252 lua_setglobal(L, "WebSocketMessageType");
253 setfield_integer(L, "TEXT", (int)ix::WebSocketMessageType::Message);
254 setfield_integer(L, "BINARY", MESSAGE_TYPE_BINARY);
255 setfield_integer(L, "OPEN", (int)ix::WebSocketMessageType::Open);
256 setfield_integer(L, "CLOSE", (int)ix::WebSocketMessageType::Close);
257 setfield_integer(L, "ERROR", (int)ix::WebSocketMessageType::Error);
258 setfield_integer(L, "PING", (int)ix::WebSocketMessageType::Ping);
259 setfield_integer(L, "PONG", (int)ix::WebSocketMessageType::Pong);
260 setfield_integer(L, "FRAGMENT", (int)ix::WebSocketMessageType::Fragment);
261 lua_pop(L, 1);
262}
263
264} // namespace script
265} // namespace app
266