Enhance WebSocket support: add opcode handling, binary message support, and improve connection validation

src/fastcgi/src/fcgicc.cc

@@ -60,6 +60,20 @@ make_http_text_response(String status_line, String body, String extra_headers =
 	);
 }
 
+static bool
+is_valid_close_code(u16 status_code)
+{
+	if(status_code < 1000)
+		return(false);
+	if(status_code == 1004 || status_code == 1005 || status_code == 1006 || status_code == 1015)
+		return(false);
+	if(status_code <= 1014)
+		return(true);
+	if(status_code >= 3000 && status_code <= 4999)
+		return(true);
+	return(false);
+}
+
 void
 FastCGIServer::shutdown()
 {
@@ -217,6 +231,36 @@ FastCGIServer::close_http_listeners()
 	}
 }
 
+void
+FastCGIServer::close_websocket_connection(Connection& connection, u16 status_code, String reason)
+{
+	if(!connection.close_socket)
+		connection.output_buffer += ws_close_frame(status_code, reason);
+	connection.close_socket = true;
+}
+
+void
+FastCGIServer::fail_websocket_connection(Connection& connection, u16 status_code, String reason)
+{
+	close_websocket_connection(connection, status_code, reason);
+}
+
+void
+FastCGIServer::dispatch_websocket_message(Connection& connection, RequestID request_id, String payload, u8 opcode)
+{
+	RequestList::iterator it = connection.requests.find(request_id);
+	if(it == connection.requests.end() || !on_websocket_message)
+		return;
+
+	it->second->resources.is_websocket = true;
+	it->second->resources.websocket_connection_id = connection.websocket_connection_id;
+	it->second->resources.websocket_scope = connection.websocket_scope;
+	it->second->resources.websocket_opcode = opcode;
+	it->second->resources.websocket_is_binary = (opcode == 0x2);
+	it->second->resources.websocket_is_text = (opcode == 0x1);
+	on_websocket_message(*it->second, payload, opcode);
+}
+
 void
 FastCGIServer::process(int timeout_ms)
 {
@@ -394,9 +438,39 @@ FastCGIServer::process_http_request(FastCGIRequest& request, String& data)
 	if(to_lower(request.params["HTTP_UPGRADE"]) == "websocket")
 	{
 		Connection* connection = client_sockets[request.resources.client_socket];
+		String request_method = trim(request.params["REQUEST_METHOD"]);
+		String request_protocol = trim(request.params["SERVER_PROTOCOL"]);
 		String websocket_connection = to_lower(request.params["HTTP_CONNECTION"]);
+		String websocket_host = trim(request.params["HTTP_HOST"]);
 		String websocket_key = trim(request.params["HTTP_SEC_WEBSOCKET_KEY"]);
 		String websocket_version = trim(request.params["HTTP_SEC_WEBSOCKET_VERSION"]);
+		if(request_method != "GET")
+		{
+			connection->output_buffer += make_http_text_response(
+				"HTTP/1.1 405 Method Not Allowed",
+				"websocket upgrades require GET"
+			);
+			connection->close_socket = true;
+			return;
+		}
+		if(request_protocol != "HTTP/1.1")
+		{
+			connection->output_buffer += make_http_text_response(
+				"HTTP/1.1 400 Bad Request",
+				"websocket upgrades require HTTP/1.1"
+			);
+			connection->close_socket = true;
+			return;
+		}
+		if(websocket_host == "")
+		{
+			connection->output_buffer += make_http_text_response(
+				"HTTP/1.1 400 Bad Request",
+				"missing Host header"
+			);
+			connection->close_socket = true;
+			return;
+		}
 		if(websocket_connection.find("upgrade") == String::npos)
 		{
 			connection->output_buffer += make_http_text_response(
@@ -415,6 +489,24 @@ FastCGIServer::process_http_request(FastCGIRequest& request, String& data)
 			connection->close_socket = true;
 			return;
 		}
+		if(!ws_is_valid_client_key(websocket_key))
+		{
+			connection->output_buffer += make_http_text_response(
+				"HTTP/1.1 400 Bad Request",
+				"invalid Sec-WebSocket-Key header"
+			);
+			connection->close_socket = true;
+			return;
+		}
+		if(websocket_version == "")
+		{
+			connection->output_buffer += make_http_text_response(
+				"HTTP/1.1 400 Bad Request",
+				"missing Sec-WebSocket-Version header"
+			);
+			connection->close_socket = true;
+			return;
+		}
 		if(websocket_version != "" && websocket_version != "13")
 		{
 			connection->output_buffer += make_http_text_response(
@@ -482,44 +574,108 @@ FastCGIServer::process_websocket_input(Connection& connection)
 
 		if(!frame.mask_bit)
 		{
-			connection.output_buffer += ws_close_frame(1002, "client frames must be masked");
-			connection.close_socket = true;
+			fail_websocket_connection(connection, 1002, "client frames must be masked");
 			return;
 		}
 
-		if(!frame.is_final_fragment || frame.opcode == 0x0)
+		bool is_control_frame = (frame.opcode & 0x08) != 0;
+		if(is_control_frame && !frame.is_final_fragment)
 		{
-			connection.output_buffer += ws_close_frame(1003, "fragmented frames are not supported");
-			connection.close_socket = true;
+			fail_websocket_connection(connection, 1002, "control frames must not be fragmented");
 			return;
 		}
 
 		switch(frame.opcode)
 		{
-			case 0x1:
+			case 0x0:
 			{
-				RequestList::iterator it = connection.requests.find(connection.client_socket);
-				if(it != connection.requests.end() && on_websocket_message)
+				if(connection.websocket_fragment_opcode == 0)
 				{
-					it->second->resources.is_websocket = true;
-					it->second->resources.websocket_connection_id = connection.websocket_connection_id;
-					it->second->resources.websocket_scope = connection.websocket_scope;
-					on_websocket_message(*it->second, frame.payload);
+					fail_websocket_connection(connection, 1002, "unexpected continuation frame");
+					return;
 				}
+
+				connection.websocket_fragment_buffer += frame.payload;
+				if(!frame.is_final_fragment)
+					break;
+
+				String payload = connection.websocket_fragment_buffer;
+				u8 opcode = connection.websocket_fragment_opcode;
+				connection.websocket_fragment_buffer = "";
+				connection.websocket_fragment_opcode = 0;
+
+				if(opcode == 0x1 && !ws_is_valid_utf8(payload))
+				{
+					fail_websocket_connection(connection, 1007, "invalid UTF-8 text message");
+					return;
+				}
+
+				dispatch_websocket_message(connection, connection.client_socket, payload, opcode);
+				break;
+			}
+			case 0x1:
+			case 0x2:
+			{
+				if(connection.websocket_fragment_opcode != 0)
+				{
+					fail_websocket_connection(connection, 1002, "new data frame while fragmented message is active");
+					return;
+				}
+
+				if(frame.is_final_fragment)
+				{
+					if(frame.opcode == 0x1 && !ws_is_valid_utf8(frame.payload))
+					{
+						fail_websocket_connection(connection, 1007, "invalid UTF-8 text message");
+						return;
+					}
+
+					dispatch_websocket_message(connection, connection.client_socket, frame.payload, frame.opcode);
+					break;
+				}
+
+				connection.websocket_fragment_buffer = frame.payload;
+				connection.websocket_fragment_opcode = frame.opcode;
 				break;
 			}
 			case 0x8:
-				connection.output_buffer += ws_close_frame();
-				connection.close_socket = true;
+			{
+				if(frame.payload.length() == 1)
+				{
+					fail_websocket_connection(connection, 1002, "invalid close frame payload");
+					return;
+				}
+
+				u16 status_code = 1000;
+				String reason = "";
+				if(frame.payload.length() >= 2)
+				{
+					status_code =
+						((u16)(u8)frame.payload[0] << 8) |
+						(u16)(u8)frame.payload[1];
+					reason = frame.payload.substr(2);
+					if(!is_valid_close_code(status_code))
+					{
+						fail_websocket_connection(connection, 1002, "invalid close status code");
+						return;
+					}
+					if(!ws_is_valid_utf8(reason))
+					{
+						fail_websocket_connection(connection, 1007, "invalid UTF-8 close reason");
+						return;
+					}
+				}
+
+				close_websocket_connection(connection, status_code, reason);
 				return;
+			}
 			case 0x9:
 				connection.output_buffer += ws_encode_frame(frame.payload, 0xA);
 				break;
 			case 0xA:
 				break;
 			default:
-				connection.output_buffer += ws_close_frame(1003, "unsupported websocket opcode");
-				connection.close_socket = true;
+				fail_websocket_connection(connection, 1002, "unsupported websocket opcode");
 				return;
 		}
 	}

src/fastcgi/src/fcgicc.h

@@ -48,7 +48,7 @@ public:
 	std::function<int(FastCGIRequest&)> on_request = 0;
 	std::function<int(FastCGIRequest&)> on_data = 0;
 	std::function<int(FastCGIRequest&)> on_complete = 0;
-	std::function<int(FastCGIRequest&, const String&)> on_websocket_message = 0;
+	std::function<int(FastCGIRequest&, const String&, u8)> on_websocket_message = 0;
 
 	int listen(unsigned tcp_port);
 	int listen_http(unsigned tcp_port);
@@ -72,6 +72,8 @@ public:
 		bool is_websocket = false;
 		String websocket_connection_id;
 		String websocket_scope;
+		String websocket_fragment_buffer;
+		u8 websocket_fragment_opcode = 0;
 		char type = 'F'; // F = FastCGI, H = HttpServer
 	};
 
@@ -86,6 +88,9 @@ public:
 	void close_http_listeners();
 	void read_fgci(Connection&);
 	void process_websocket_input(Connection&);
+	void fail_websocket_connection(Connection& connection, u16 status_code, String reason = "");
+	void close_websocket_connection(Connection& connection, u16 status_code = 1000, String reason = "");
+	void dispatch_websocket_message(Connection& connection, RequestID request_id, String payload, u8 opcode);
 	bool websocket_send_to(String connection_id, String message);
 	u64 websocket_broadcast(String scope, String message);
 	StringList websocket_connection_ids(String scope = "");

src/lib/sys.h

@@ -39,6 +39,8 @@ String socket_read(u64 sockfd, u32 max_length = 1024*128, u32 timeout = 1);
 String ws_message();
 String ws_connection_id();
 String ws_scope();
+u8 ws_opcode();
+bool ws_is_binary();
 StringList ws_connections(String scope = "");
 u64 ws_connection_count(String scope = "");
 bool ws_send(String message, String scope = "");

src/lib/types.h

@@ -173,6 +173,9 @@ struct Request {
 		bool is_websocket = false;
 		String websocket_connection_id = "";
 		String websocket_scope = "";
+		u8 websocket_opcode = 0;
+		bool websocket_is_binary = false;
+		bool websocket_is_text = false;
 		std::string params_buffer;
 	} resources;
 

src/lib/uri.cpp

@@ -27,6 +27,134 @@ static String base64_encode(String raw)
 	return(result);
 }
 
+static int base64_decode_value(char c)
+{
+	if(c >= 'A' && c <= 'Z')
+		return(c - 'A');
+	if(c >= 'a' && c <= 'z')
+		return(c - 'a' + 26);
+	if(c >= '0' && c <= '9')
+		return(c - '0' + 52);
+	if(c == '+')
+		return(62);
+	if(c == '/')
+		return(63);
+	return(-1);
+}
+
+static String base64_decode(String raw, bool& ok)
+{
+	ok = false;
+	String cleaned;
+	for(char c : raw)
+	{
+		if(!isspace(c))
+			cleaned.append(1, c);
+	}
+
+	if(cleaned.length() == 0 || (cleaned.length() % 4) != 0)
+		return("");
+
+	String result;
+	for(u32 i = 0; i < cleaned.length(); i += 4)
+	{
+		int values[4];
+		int padding = 0;
+		for(u32 j = 0; j < 4; j++)
+		{
+			char c = cleaned[i + j];
+			if(c == '=')
+			{
+				values[j] = 0;
+				padding += 1;
+			}
+			else
+			{
+				values[j] = base64_decode_value(c);
+				if(values[j] < 0)
+					return("");
+			}
+		}
+
+		if(padding > 2)
+			return("");
+		if(padding > 0 && i + 4 != cleaned.length())
+			return("");
+		if(cleaned[i + 2] == '=' && cleaned[i + 3] != '=')
+			return("");
+
+		result.append(1, (char)((values[0] << 2) | (values[1] >> 4)));
+		if(cleaned[i + 2] != '=')
+			result.append(1, (char)(((values[1] & 0x0F) << 4) | (values[2] >> 2)));
+		if(cleaned[i + 3] != '=')
+			result.append(1, (char)(((values[2] & 0x03) << 6) | values[3]));
+	}
+
+	ok = true;
+	return(result);
+}
+
+bool ws_is_valid_utf8(String input)
+{
+	u32 i = 0;
+	while(i < input.length())
+	{
+		u8 c = (u8)input[i];
+		u32 trailing = 0;
+		u32 codepoint = 0;
+
+		if(c <= 0x7F)
+		{
+			i += 1;
+			continue;
+		}
+		else if((c & 0xE0) == 0xC0)
+		{
+			trailing = 1;
+			codepoint = c & 0x1F;
+			if(codepoint == 0)
+				return(false);
+		}
+		else if((c & 0xF0) == 0xE0)
+		{
+			trailing = 2;
+			codepoint = c & 0x0F;
+		}
+		else if((c & 0xF8) == 0xF0)
+		{
+			trailing = 3;
+			codepoint = c & 0x07;
+		}
+		else
+		{
+			return(false);
+		}
+
+		if(i + trailing >= input.length())
+			return(false);
+
+		for(u32 j = 1; j <= trailing; j++)
+		{
+			u8 follow = (u8)input[i + j];
+			if((follow & 0xC0) != 0x80)
+				return(false);
+			codepoint = (codepoint << 6) | (follow & 0x3F);
+		}
+
+		if((trailing == 1 && codepoint < 0x80) ||
+			(trailing == 2 && codepoint < 0x800) ||
+			(trailing == 3 && codepoint < 0x10000))
+			return(false);
+		if(codepoint > 0x10FFFF)
+			return(false);
+		if(codepoint >= 0xD800 && codepoint <= 0xDFFF)
+			return(false);
+
+		i += trailing + 1;
+	}
+	return(true);
+}
+
 String var_dump(URI uri, String prefix, String postfix)
 {
 	return(
@@ -463,6 +591,13 @@ String ws_make_accept_key(String client_key)
 	)));
 }
 
+bool ws_is_valid_client_key(String client_key)
+{
+	bool ok = false;
+	String decoded = base64_decode(trim(client_key), ok);
+	return(ok && decoded.length() == 16);
+}
+
 String ws_encode_frame(String payload, u8 opcode, bool is_final_fragment)
 {
 	String frame;
@@ -513,6 +648,9 @@ bool WSFrame::parse(const String& buffer, String& error)
 	const unsigned char* raw = (const unsigned char*)buffer.data();
 	opcode = raw[0] & 0x0F;
 	is_final_fragment = (raw[0] & 0x80) != 0;
+	rsv1 = (raw[0] & 0x40) != 0;
+	rsv2 = (raw[0] & 0x20) != 0;
+	rsv3 = (raw[0] & 0x10) != 0;
 	mask_bit = (raw[1] & 0x80) != 0;
 	payload_length = raw[1] & 0x7F;
 	header_length = 2;
@@ -528,6 +666,11 @@ bool WSFrame::parse(const String& buffer, String& error)
 	{
 		if(buffer.length() < 10)
 			return(false);
+		if((raw[2] & 0x80) != 0)
+		{
+			error = "invalid websocket frame length";
+			return(false);
+		}
 		payload_length = 0;
 		for(u32 i = 0; i < 8; i++)
 			payload_length = (payload_length << 8) | (u64)raw[2 + i];
@@ -544,6 +687,17 @@ bool WSFrame::parse(const String& buffer, String& error)
 		error = "invalid websocket frame length";
 		return(false);
 	}
+	if(rsv1 || rsv2 || rsv3)
+	{
+		error = "reserved websocket bits are not supported";
+		return(false);
+	}
+	bool is_control_frame = (opcode & 0x08) != 0;
+	if(is_control_frame && payload_length > 125)
+	{
+		error = "control frames must be 125 bytes or less";
+		return(false);
+	}
 	if(buffer.length() < frame_length)
 		return(false);
 

src/lib/uri.h

@@ -18,14 +18,19 @@ void save_session_data(String session_id, StringMap data);
 String session_start(String session_name = "uce-session");
 void session_destroy(String session_name = "uce-session");
 String ws_make_accept_key(String client_key);
+bool ws_is_valid_client_key(String client_key);
 String ws_encode_frame(String payload, u8 opcode = 0x1, bool is_final_fragment = true);
 String ws_close_frame(u16 status_code = 1000, String reason = "");
+bool ws_is_valid_utf8(String input);
 
 struct WSFrame {
 
 	u8 opcode = 0;
 	bool is_final_fragment = false;
 	bool mask_bit = false;
+	bool rsv1 = false;
+	bool rsv2 = false;
+	bool rsv3 = false;
 	u64 payload_length = 0;
 	u64 header_length = 0;
 	u64 frame_length = 0;

src/linux_fastcgi.cpp

@@ -58,6 +58,20 @@ String ws_scope()
 	return(current_ws_scope());
 }
 
+u8 ws_opcode()
+{
+	if(!context)
+		return(0);
+	return(context->resources.websocket_opcode);
+}
+
+bool ws_is_binary()
+{
+	if(!context)
+		return(false);
+	return(context->resources.websocket_is_binary);
+}
+
 StringList ws_connections(String scope)
 {
 	return(server.websocket_connection_ids(normalize_ws_scope(scope)));
@@ -173,7 +187,7 @@ int handle_complete(FastCGIRequest& request) {
     return 0;
 }
 
-int handle_websocket_message(FastCGIRequest& request, const String& message)
+int handle_websocket_message(FastCGIRequest& request, const String& message, u8 opcode)
 {
 	Request event_request;
 	ByteStream ws_output;
@@ -202,6 +216,9 @@ int handle_websocket_message(FastCGIRequest& request, const String& message)
 	event_request.var["ws"]["connection_id"] = request.resources.websocket_connection_id;
 	event_request.var["ws"]["scope"] = request.resources.websocket_scope;
 	event_request.var["ws"]["connection_count"] = (f64)server.websocket_connection_ids(request.resources.websocket_scope).size();
+	event_request.var["ws"]["opcode"] = (f64)opcode;
+	event_request.var["ws"]["is_binary"].set_bool(request.resources.websocket_is_binary);
+	event_request.var["ws"]["is_text"].set_bool(request.resources.websocket_is_text);
 	event_request.var["ws"]["document_uri"] = first(
 		request.params["DOCUMENT_URI"],
 		request.params["REQUEST_URI"]
@@ -210,6 +227,7 @@ int handle_websocket_message(FastCGIRequest& request, const String& message)
 	call_param["message"] = message;
 	call_param["connection_id"] = request.resources.websocket_connection_id;
 	call_param["scope"] = request.resources.websocket_scope;
+	call_param["opcode"] = (f64)opcode;
 	call_param["document_uri"] = event_request.var["ws"]["document_uri"].to_string();
 
 	compiler_invoke_websocket(&event_request, request.params["SCRIPT_FILENAME"], call_param);