From 5157e60379dbd5796f27a3bc089112465c8bd377 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Sun, 12 Feb 2023 18:24:55 +0100
Subject: [PATCH] Improve websocket examples

---
 .lune/examples/websocket_client.luau  | 11 ++++++++---
 .lune/examples/websocket_server.luau  |  5 +++--
 luneTypes.d.luau                      |  1 +
 packages/lib/src/lua/net/ws_client.rs | 14 ++++----------
 packages/lib/src/lua/net/ws_server.rs | 14 ++++----------
 5 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/.lune/examples/websocket_client.luau b/.lune/examples/websocket_client.luau
index 5be1b8a..5fa015b 100644
--- a/.lune/examples/websocket_client.luau
+++ b/.lune/examples/websocket_client.luau
@@ -1,6 +1,10 @@
---> A basic web socket client communicates with an echo server
+--> A basic web socket client that communicates with an echo server
 
-local URL = "wss://demo.piesocket.com/v3/"
+local PORT = if process.env.PORT ~= nil and #process.env.PORT > 0
+	then assert(tonumber(process.env.PORT), "Failed to parse port from env")
+	else 8080
+
+local URL = `ws://127.0.0.1:{PORT}`
 
 -- Connect to our web socket server
 
@@ -11,7 +15,7 @@ print("Sending a message every second for 5 seconds...")
 
 -- Force exit after 10 seconds in case the server is not responding well
 
-task.delay(10, function()
+local forceExit = task.delay(10, function()
 	warn("Example did not complete in time, exiting...")
 	process.exit(1)
 end)
@@ -32,4 +36,5 @@ end
 print("Closing web socket...")
 socket.close()
 
+task.cancel(forceExit)
 print("Done! 🌙")
diff --git a/.lune/examples/websocket_server.luau b/.lune/examples/websocket_server.luau
index 66d0f93..3e85c9e 100644
--- a/.lune/examples/websocket_server.luau
+++ b/.lune/examples/websocket_server.luau
@@ -13,7 +13,7 @@ local handle = net.serve(PORT, {
 		repeat
 			local message = socket.next()
 			if message ~= nil then
-				socket.send("Echo\n" .. message)
+				socket.send("Echo - " .. message)
 			end
 		until message == nil
 		print("Web socket disconnected.")
@@ -25,8 +25,9 @@ print(`Listening on port {PORT} 🚀`)
 -- Exit our example after a small delay, if you copy this
 -- example just remove this part to keep the server running
 
-task.delay(2, function()
+task.delay(10, function()
 	print("Shutting down...")
 	task.wait(1)
 	handle.stop()
+	task.wait(1)
 end)
diff --git a/luneTypes.d.luau b/luneTypes.d.luau
index 940e7d1..4ba5d73 100644
--- a/luneTypes.d.luau
+++ b/luneTypes.d.luau
@@ -169,6 +169,7 @@ declare class NetWebSocket
 	close: () -> ()
 	send: (message: string) -> ()
 	next: () -> string?
+	iter: () -> () -> string
 	function __iter(self): () -> string
 end
 
diff --git a/packages/lib/src/lua/net/ws_client.rs b/packages/lib/src/lua/net/ws_client.rs
index 3cdf578..c27b423 100644
--- a/packages/lib/src/lua/net/ws_client.rs
+++ b/packages/lib/src/lua/net/ws_client.rs
@@ -39,15 +39,9 @@ impl NetWebSocketClient {
         local proxy = newproxy(true)
         local meta = getmetatable(proxy)
         meta.__index = {
-            close = function()
-                return ws:close()
-            end,
-            send = function(...)
-                return ws:send(...)
-            end,
-            next = function()
-                return ws:next()
-            end,
+            close = function()    return ws:close()   end,
+            send  = function(...) return ws:send(...) end,
+            next  = function()    return ws:next()    end,
         }
         meta.__iter = function()
             return function()
@@ -75,7 +69,7 @@ impl LuaUserData for NetWebSocketClient {
                 }),
                 None => Ok(LuaValue::Nil),
             }
-        })
+        });
     }
 }
 
diff --git a/packages/lib/src/lua/net/ws_server.rs b/packages/lib/src/lua/net/ws_server.rs
index 18966fa..be8b42e 100644
--- a/packages/lib/src/lua/net/ws_server.rs
+++ b/packages/lib/src/lua/net/ws_server.rs
@@ -39,15 +39,9 @@ impl NetWebSocketServer {
         local proxy = newproxy(true)
         local meta = getmetatable(proxy)
         meta.__index = {
-            close = function()
-                return ws:close()
-            end,
-            send = function(...)
-                return ws:send(...)
-            end,
-            next = function()
-                return ws:next()
-            end,
+            close = function()    return ws:close()   end,
+            send  = function(...) return ws:send(...) end,
+            next  = function()    return ws:next()    end,
         }
         meta.__iter = function()
             return function()
@@ -75,7 +69,7 @@ impl LuaUserData for NetWebSocketServer {
                 }),
                 None => Ok(LuaValue::Nil),
             }
-        })
+        });
     }
 }