From aaa0977d549189ad502154c6b2d4b960396c41ad Mon Sep 17 00:00:00 2001
From: reggie <reggie.latte@gmail.com>
Date: Thu, 15 Jun 2023 15:36:05 -0500
Subject: [PATCH] Add TLS support for `net.socket()`, fix `WebSocket.closeCode`
 behavior (#57)

---
 Cargo.lock                            | 16 +++++++++++++++-
 packages/lib/Cargo.toml               |  2 +-
 packages/lib/src/lua/net/websocket.rs |  2 +-
 packages/lib/src/tests.rs             |  1 +
 tests/net/socket/wss.luau             | 25 +++++++++++++++++++++++++
 5 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 tests/net/socket/wss.luau

diff --git a/Cargo.lock b/Cargo.lock
index 01ab725..d72512a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1749,7 +1749,7 @@ dependencies = [
  "wasm-bindgen",
  "wasm-bindgen-futures",
  "web-sys",
- "webpki-roots",
+ "webpki-roots 0.22.6",
  "winreg",
 ]
 
@@ -2361,8 +2361,11 @@ checksum = "ec509ac96e9a0c43427c74f003127d953a265737636129424288d27cb5c4b12c"
 dependencies = [
  "futures-util",
  "log",
+ "rustls",
  "tokio",
+ "tokio-rustls",
  "tungstenite",
+ "webpki-roots 0.23.1",
 ]
 
 [[package]]
@@ -2459,10 +2462,12 @@ dependencies = [
  "httparse",
  "log",
  "rand",
+ "rustls",
  "sha1 0.10.5",
  "thiserror",
  "url",
  "utf-8",
+ "webpki",
 ]
 
 [[package]]
@@ -2684,6 +2689,15 @@ dependencies = [
  "webpki",
 ]
 
+[[package]]
+name = "webpki-roots"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b03058f88386e5ff5310d9111d53f48b17d732b401aeb83a8d5190f2ac459338"
+dependencies = [
+ "rustls-webpki",
+]
+
 [[package]]
 name = "winapi"
 version = "0.3.9"
diff --git a/packages/lib/Cargo.toml b/packages/lib/Cargo.toml
index 4b3098c..6d96df3 100644
--- a/packages/lib/Cargo.toml
+++ b/packages/lib/Cargo.toml
@@ -52,7 +52,7 @@ async-compression = { version = "0.4", features = [
 ] }
 hyper = { version = "0.14", features = ["full"] }
 hyper-tungstenite = { version = "0.10" }
-tokio-tungstenite = { version = "0.19" }
+tokio-tungstenite = { version = "0.19", features = ["rustls-tls-webpki-roots"] }
 
 [dev-dependencies]
 anyhow = "1.0"
diff --git a/packages/lib/src/lua/net/websocket.rs b/packages/lib/src/lua/net/websocket.rs
index e1ce2b9..bbf24e0 100644
--- a/packages/lib/src/lua/net/websocket.rs
+++ b/packages/lib/src/lua/net/websocket.rs
@@ -35,7 +35,7 @@ return freeze(setmetatable({
 }, {
 	__index = function(self, key)
 		if key == "closeCode" then
-			return close_code()
+			return close_code(websocket)
 		end
 	end,
 }))
diff --git a/packages/lib/src/tests.rs b/packages/lib/src/tests.rs
index 7d6abd6..e398f27 100644
--- a/packages/lib/src/tests.rs
+++ b/packages/lib/src/tests.rs
@@ -57,6 +57,7 @@ create_tests! {
     net_url_decode: "net/url/decode",
     net_serve_requests: "net/serve/requests",
     net_serve_websockets: "net/serve/websockets",
+    net_socket_wss: "net/socket/wss",
 
     process_args: "process/args",
     process_cwd: "process/cwd",
diff --git a/tests/net/socket/wss.luau b/tests/net/socket/wss.luau
new file mode 100644
index 0000000..a8a841c
--- /dev/null
+++ b/tests/net/socket/wss.luau
@@ -0,0 +1,25 @@
+local net = require("@lune/net")
+local stdio = require("@lune/stdio")
+local task = require("@lune/task")
+
+-- net.socket() will inherently spit out an error if init connection fails, but for
+-- the sake of this test, prot. call checks are unnecessary. Also, we're going to use
+-- Discord's WebSocket gateway server for soley testing wss, as we don't need to auth
+-- or anything for that; we just want to make sure it functions..
+print("Connecting to Discord wss gateway..")
+local socket = net.socket("wss://gateway.discord.gg/?v=10&encoding=json")
+print("Connected!") -- Would've errored at net.socket() call before this if init connection err
+
+while not socket.closeCode do
+    local nextMessage = socket.next()
+
+    if nextMessage then
+        print(`{stdio.style("bold")}Discord:{stdio.style("reset")}`, nextMessage)
+        
+        print("Will (deliberately) close the socket in 2 seconds..")
+        task.wait(2)
+        socket.close(1000)
+    end
+end
+
+print(`Connection to socket closed with closeCode {socket.closeCode}`)