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

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;