First working draft of a websocket binding for Lua

src/civetweb.c

@@ -518,6 +518,11 @@ enum {
     GLOBAL_PASSWORDS_FILE, INDEX_FILES, ENABLE_KEEP_ALIVE, ACCESS_CONTROL_LIST,
     EXTRA_MIME_TYPES, LISTENING_PORTS, DOCUMENT_ROOT, SSL_CERTIFICATE,
     NUM_THREADS, RUN_AS_USER, REWRITE, HIDE_FILES, REQUEST_TIMEOUT,
+
+#if defined(USE_LUA) && defined(USE_WEBSOCKET)
+    LUA_WEBSOCKET_SCRIPT,
+#endif
+
     NUM_OPTIONS
 };
 
@@ -547,6 +552,11 @@ static const char *config_options[] = {
     "url_rewrite_patterns", NULL,
     "hide_files_patterns", NULL,
     "request_timeout_ms", "30000",
+
+#if defined(USE_LUA) && defined(USE_WEBSOCKET)
+    "lua_websocket_script", NULL,
+#endif
+
     NULL
 };
 
@@ -610,6 +620,9 @@ struct mg_connection {
     int64_t last_throttle_bytes;/* Bytes sent this second */
     pthread_mutex_t mutex;      /* Used by mg_lock/mg_unlock to ensure atomic
                                    transmissions for websockets */
+#if defined(USE_LUA) && defined(USE_WEBSOCKET)
+    void * lua_websocket_state; /* Lua_State for a websocket connection */
+#endif
 };
 
 static pthread_key_t sTlsKey;  /* Thread local storage index */
@@ -4244,6 +4257,10 @@ void mg_unlock(struct mg_connection* conn)
     (void) pthread_mutex_unlock(&conn->mutex);
 }
 
+#ifdef USE_LUA
+#include "mod_lua.inl"
+#endif /* USE_LUA */
+
 #if defined(USE_WEBSOCKET)
 
 /* START OF SHA-1 code
@@ -4597,6 +4614,10 @@ static void read_websocket(struct mg_connection *conn)
                or "connection close" opcode received. */
             if ((conn->ctx->callbacks.websocket_data != NULL &&
                  !conn->ctx->callbacks.websocket_data(conn, mop, data, data_len)) ||
+#ifdef USE_LUA
+                (conn->lua_websocket_state &&
+                 !lua_websocket_data(conn, mop, data, data_len)) ||
+#endif
                 (buf[0] & 0xf) == WEBSOCKET_OPCODE_CONNECTION_CLOSE) {  /* Opcode == 8, connection close */
                 break;
             }
@@ -4660,18 +4681,34 @@ int mg_websocket_write(struct mg_connection* conn, int opcode, const char* data,
 
 static void handle_websocket_request(struct mg_connection *conn)
 {
+    const char *ws_page = NULL;
     const char *version = mg_get_header(conn, "Sec-WebSocket-Version");
     if (version == NULL || strcmp(version, "13") != 0) {
         send_http_error(conn, 426, "Upgrade Required", "%s", "Upgrade Required");
     } else if (conn->ctx->callbacks.websocket_connect != NULL &&
                conn->ctx->callbacks.websocket_connect(conn) != 0) {
-        /* Callback has returned non-zero, do not proceed with handshake */
+        /* C callback has returned non-zero, do not proceed with handshake. */
+        /* The C callback is called before Lua and may prevent Lua from handling the websocket. */
     } else {
-        send_websocket_handshake(conn);
-        if (conn->ctx->callbacks.websocket_ready != NULL) {
-            conn->ctx->callbacks.websocket_ready(conn);
+#ifdef USE_LUA
+        ws_page = conn->ctx->config[LUA_WEBSOCKET_SCRIPT];
+        if (ws_page) {
+            conn->lua_websocket_state = new_lua_websocket(ws_page, conn);
+            if (conn->lua_websocket_state) {
+                send_websocket_handshake(conn);
+                lua_websocket_ready(conn);
+                read_websocket(conn);
+            }
+        } else
+#endif
+        {
+            /* No Lua websock script specified. */
+            send_websocket_handshake(conn);
+            if (conn->ctx->callbacks.websocket_ready != NULL) {
+                conn->ctx->callbacks.websocket_ready(conn);
+            }
+            read_websocket(conn);
         }
-        read_websocket(conn);
     }
 }
 
@@ -4747,10 +4784,6 @@ static uint32_t get_remote_ip(const struct mg_connection *conn)
     return ntohl(* (uint32_t *) &conn->client.rsa.sin.sin_addr);
 }
 
-#ifdef USE_LUA
-#include "mod_lua.inl"
-#endif /* USE_LUA */
-
 int mg_upload(struct mg_connection *conn, const char *destination_dir)
 {
     const char *content_type_header, *boundary_start;
@@ -5560,6 +5593,12 @@ static void close_socket_gracefully(struct mg_connection *conn)
 
 static void close_connection(struct mg_connection *conn)
 {
+#if defined(USE_LUA) && defined(USE_WEBSOCKET)
+    if (conn->lua_websocket_state) {
+        lua_websocket_close(conn);
+    }
+#endif
+
     /* call the connection_close callback if assigned */
     if (conn->ctx->callbacks.connection_close != NULL)
         conn->ctx->callbacks.connection_close(conn);

src/mod_lua.inl

@@ -11,7 +11,12 @@ static void *mmap(void *addr, int64_t len, int prot, int flags, int fd,
     CloseHandle(mh);
     return p;
 }
-#define munmap(x, y)  UnmapViewOfFile(x)
+
+static void munmap(void *addr, int64_t length)
+{
+    UnmapViewOfFile(addr);
+}
+
 #define MAP_FAILED NULL
 #define MAP_PRIVATE 0
 #define PROT_READ 0
@@ -424,3 +429,112 @@ static int handle_lsp_request(struct mg_connection *conn, const char *path,
     conn->must_close=1;
     return error;
 }
+
+static void * new_lua_websocket(const char * script, struct mg_connection *conn)
+{
+    lua_State *L = NULL;
+    int ok = 0;
+    int err;
+
+    assert(conn->lua_websocket_state == NULL);
+    L = luaL_newstate();
+    if (L) {
+        prepare_lua_environment(conn, L);
+        if (conn->ctx->callbacks.init_lua != NULL) {
+            conn->ctx->callbacks.init_lua(conn, L);
+        }
+        err = luaL_loadfile(L, script);
+        switch (err) {
+            case 0:
+                {
+                    err = lua_pcall(L, 0, LUA_MULTRET, 0);
+                    switch (err) {
+                        case 0:
+                            /* return nothing or true to continue, false to stop */
+                            ok = !lua_isboolean(L, -1) || lua_toboolean(L, -1);
+                            break;
+                        case LUA_ERRMEM:
+                            mg_cry(conn, "%s: lua_pcall failed: out of memory", __func__);
+                            break;
+                        case LUA_ERRRUN:
+                            mg_cry(conn, "%s: lua_pcall failed: runtime error: %s", __func__, lua_tostring(L, -1));
+                            break;
+                        case LUA_ERRERR:
+                            mg_cry(conn, "%s: lua_pcall failed: double fault: %s", __func__, lua_tostring(L, -1));
+                            break;
+                        default:
+                            mg_cry(conn, "%s: lua_pcall failed: error %i", __func__, err);
+                            break;
+                    }
+                }
+                break;
+            case LUA_ERRMEM:
+                mg_cry(conn, "%s: luaL_loadfile failed: out of memory", __func__);
+                break;
+            case LUA_ERRFILE:
+                mg_cry(conn, "%s: luaL_loadfile failed: file %s not found", __func__, script);
+                break;
+            case LUA_ERRSYNTAX:
+                mg_cry(conn, "%s: luaL_loadfile failed: syntax error: %s", __func__, lua_tostring(L, -1));
+                break;
+            default:
+                mg_cry(conn, "%s: luaL_loadfile failed: error %i", __func__, err);
+                break;
+        }
+        if (!ok) {
+            lua_close(L);
+            L = NULL;
+        }
+    } else {
+        mg_cry(conn, "%s: luaL_newstate failed", __func__);
+    }
+
+    return L;
+}
+
+static void lua_websocket_ready(struct mg_connection *conn)
+{
+    lua_State *L = (lua_State*)(conn->lua_websocket_state);
+
+    assert(L != NULL);
+
+    lua_getglobal(L, "ready");
+    if (lua_pcall(L, 0, 0, 0) != 0) {
+        mg_cry(conn, "%s: error running function `ready': %s", lua_tostring(L, -1));
+    }
+}
+
+static int lua_websocket_data(struct mg_connection *conn, int bits, char *data, size_t data_len)
+{
+    lua_State *L = (lua_State*)(conn->lua_websocket_state);
+    int ok = 0;
+
+    assert(L != NULL);
+
+    lua_getglobal(L, "data");
+    lua_pushinteger(L, bits);
+    lua_pushlstring(L, data, data_len);
+    if (lua_pcall(L, 2, 1, 0) != 0) {
+        mg_cry(conn, "%s: error running function `data': %s", lua_tostring(L, -1));
+    } else {
+        ok = lua_isboolean(L, -1) && lua_toboolean(L, -1);
+    }
+
+    return ok;
+}
+
+
+static void lua_websocket_close(struct mg_connection *conn)
+{
+    lua_State *L = (lua_State*)(conn->lua_websocket_state);
+
+    assert(L != NULL);
+
+    lua_getglobal(L, "close");
+    if (lua_pcall(L, 0, 0, 0) != 0) {
+        mg_cry(conn, "%s: error running function `close': %s", lua_tostring(L, -1));
+    }
+
+    lua_close(L);
+    conn->lua_websocket_state = NULL;
+}

test/websocket.lua

@@ -0,0 +1,53 @@
+-- Open database
+local db = sqlite3.open('requests.db')
+
+if db then
+  -- Create a table if it is not created already
+  db:exec([[
+    CREATE TABLE IF NOT EXISTS requests (
+      id INTEGER PRIMARY KEY AUTOINCREMENT,
+      timestamp NOT NULL,
+      method NOT NULL,
+      uri NOT NULL,
+      addr
+    );
+  ]])
+end
+
+local function logDB(method)
+  -- Add entry about this request
+  local r;
+  repeat
+    r = db:exec([[INSERT INTO requests VALUES(NULL, datetime("now"), "]] .. method .. [[", "]] .. mg.request_info.uri .. [[", "]] .. mg.request_info.remote_port .. [[");]]);
+  until r~=5;
+
+  --[[
+  -- alternative logging (to a file)
+  local f = io.open("R:\\log.txt", "a");
+  f:write(os.date() .. " - " .. method .. " - " .. mg.request_info.uri .. " - " .. mg.request_info.remote_port .. " <" .. r .. ">\n")
+  f:close()
+  --]]
+end
+
+
+-- Callback for "Websocket ready"
+function ready()
+  logDB("WEBSOCKET READY")
+end
+
+-- Callback for "Websocket received data"
+function data(bits, content)
+    logDB(string.format("WEBSOCKET DATA (%x)", bits))
+    return true;
+end
+
+-- Callback for "Websocket is closing"
+function close()
+  logDB("WEBSOCKET CLOSE")
+  -- Close database
+  db:close()
+end
+
+
+logDB("WEBSOCKET PREPARE")
+return true; -- could return false to reject the connection before the websocket handshake