improve path handling

1 files changed, 16 insertions, 14 deletions

diff --git a/main.lua b/main.lua
index 9673e28..18fe919 100644
--- a/main.lua
+++ b/main.lua
@@ -9,28 +9,30 @@ local headers = require "http.headers"
 
 local logic = require("server_logic")
 
+local function reject(stream, code, reason)
+	local response = headers.new()
+	response:append(":status", tostring(code))
+	stream:write_headers(response, false)
+	stream:write_chunk(reason, true)
+end
+
 function onstream(server, stream)
     print("stream you.")
 	local head = stream:get_headers()
+
+	local path = head:get(":path")
+	local room = path:match '^/room/(%w+)/ws$'
+	if not room then
+		reject(stream, 404, "path not found: "..path)
+		return
+	end
+
 	local ws = websocket.new_from_stream(stream, head)
 	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)
+		reject(stream, 400, "bad request: not a websocket handshake")
 		return
 	end
 	
-	local path = head:get(":path")
-	if not path:sub(1,6) == "/room/" or not path:sub(#path-2,#path) == "/ws" or #path < 10 then
-		local response = headers.new()
-		response:append(":status", "404")
-		stream:write_headers(response, false)
-		stream:write_chunk("please access /room/(roomname).", true)
-		return
-	end
-	local room = path:sub(7,#path-3)
-
 	ws:accept()
 	logic.onsocket(ws,room)