From a7fab69838b573ae829764d6eddfa4e847e20774 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Tue, 21 Feb 2023 00:06:56 +0100
Subject: [PATCH] More websocket generics

---
 packages/lib/src/lua/net/websocket.rs | 64 +++++++++++----------------
 1 file changed, 25 insertions(+), 39 deletions(-)

diff --git a/packages/lib/src/lua/net/websocket.rs b/packages/lib/src/lua/net/websocket.rs
index d79ad8c..3098c05 100644
--- a/packages/lib/src/lua/net/websocket.rs
+++ b/packages/lib/src/lua/net/websocket.rs
@@ -1,6 +1,5 @@
 use std::sync::Arc;
 
-use hyper::upgrade::Upgraded;
 use mlua::prelude::*;
 
 use hyper_tungstenite::{tungstenite::Message as WsMessage, WebSocketStream};
@@ -8,10 +7,8 @@ use hyper_tungstenite::{tungstenite::Message as WsMessage, WebSocketStream};
 use futures_util::{SinkExt, StreamExt};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
-    net::TcpStream,
     sync::Mutex,
 };
-use tokio_tungstenite::MaybeTlsStream;
 
 use crate::utils::table::TableBuilder;
 
@@ -31,56 +28,45 @@ where
     }
 
     pub async fn close(&self) -> LuaResult<()> {
-        let mut inner = self.stream.lock().await;
-        inner.close(None).await.map_err(LuaError::external)?;
+        let mut ws = self.stream.lock().await;
+        ws.close(None).await.map_err(LuaError::external)?;
         Ok(())
     }
 
     pub async fn send(&self, msg: WsMessage) -> LuaResult<()> {
-        let mut inner = self.stream.lock().await;
-        inner.send(msg).await.map_err(LuaError::external)?;
+        let mut ws = self.stream.lock().await;
+        ws.send(msg).await.map_err(LuaError::external)?;
         Ok(())
     }
 
     pub async fn next(&self) -> LuaResult<Option<WsMessage>> {
-        let mut inner = self.stream.lock().await;
-        let item = inner.next().await.transpose();
+        let mut ws = self.stream.lock().await;
+        let item = ws.next().await.transpose();
         item.map_err(LuaError::external)
     }
-
-    pub async fn send_string(&self, msg: String) -> LuaResult<()> {
-        self.send(WsMessage::Text(msg)).await
-    }
-
-    pub async fn next_lua_value(&self, lua: &'static Lua) -> LuaResult<LuaValue> {
-        Ok(match self.next().await? {
-            Some(WsMessage::Binary(bin)) => LuaValue::String(lua.create_string(&bin)?),
-            Some(WsMessage::Text(txt)) => LuaValue::String(lua.create_string(&txt)?),
-            _ => LuaValue::Nil,
-        })
-    }
 }
 
-impl NetWebSocket<MaybeTlsStream<TcpStream>> {
+impl<T> NetWebSocket<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin + 'static,
+{
     pub fn into_lua_table(self, lua: &'static Lua) -> LuaResult<LuaTable> {
-        // FIXME: Deallocate when closed
-        let sock = Box::leak(Box::new(self));
+        let ws = Box::leak(Box::new(self));
         TableBuilder::new(lua)?
-            .with_async_function("close", |_, ()| sock.close())?
-            .with_async_function("send", |_, msg: String| sock.send_string(msg))?
-            .with_async_function("next", |lua, ()| sock.next_lua_value(lua))?
-            .build_readonly()
-    }
-}
-
-impl NetWebSocket<Upgraded> {
-    pub fn into_lua_table(self, lua: &'static Lua) -> LuaResult<LuaTable> {
-        // FIXME: Deallocate when closed
-        let sock = Box::leak(Box::new(self));
-        TableBuilder::new(lua)?
-            .with_async_function("close", |_, ()| sock.close())?
-            .with_async_function("send", |_, msg: String| sock.send_string(msg))?
-            .with_async_function("next", |lua, ()| sock.next_lua_value(lua))?
+            .with_async_function("close", |_, _: ()| async { ws.close().await })?
+            .with_async_function("send", |_, msg: String| async {
+                ws.send(WsMessage::Text(msg)).await
+            })?
+            .with_async_function("next", |_, _: ()| async {
+                match ws.next().await? {
+                    Some(msg) => Ok(match msg {
+                        WsMessage::Binary(bin) => LuaValue::String(lua.create_string(&bin)?),
+                        WsMessage::Text(txt) => LuaValue::String(lua.create_string(&txt)?),
+                        _ => LuaValue::Nil,
+                    }),
+                    None => Ok(LuaValue::Nil),
+                }
+            })?
             .build_readonly()
     }
 }