unbreak the code

1 files changed, 21 insertions, 17 deletions

diff --git a/main.lua b/main.lua
index 025ab01..cf03428 100644
--- a/main.lua
+++ b/main.lua
@@ -5,31 +5,35 @@ local cq = cqueues.new()
 
 local httpserver = require("http.server")
 local websocket = require("http.websocket")
+local headers = require "http.headers"
 
 function onstream(server, stream)
     print("stream you.")
-    cq:wrap(function()
-        local s, ws = pcall(websocket.new_from_stream,stream)
-        if not s then
-            stream:shutdown()
-            return
-        end
-        ws:accept()
-        while true do
-            local data = ws:receive()
-            ws:send(data)
-        end
-    end)
+	local ws = websocket.new_from_stream(stream, stream:get_headers())
+	if not ws then
+		local response = headers.new()
+		response:append(":status", "400")
+		stream:write_headers(response, false)
+		stream:write_chunk("bad request: not a websocket handshake", true)
+		return
+	end
+	ws:accept()
+	while true do
+		local data = ws:receive()
+		if not data then return end
+		ws:send(data)
+	end
 end
 
 local sv = httpserver.listen {
-    host = "127.0.0.1",
-    port = "61111",
-    onstream = onstream,
+	host = "127.0.0.1",
+	port = "61111",
+	onstream = onstream,
 }
 
+sv:listen()
 cq:wrap(function()
-    sv:loop()
+	assert(sv:loop())
 end)
 
-assert(cq:loop())
\ No newline at end of file
+assert(cq:loop())