adds limit to network read buffer (to prevent DoS attacks), refactors client handshake reading to leave non-handshake data in the read buffer

src/websocket_client.cpp

@@ -46,11 +46,13 @@ client::client(boost::asio::io_service& io_service,
 	  m_def_con_handler(defc) {}
 
 void client::init() {
+	// TODO: sanity check whether the session buffer size bound could be reduced
 	m_client_session = client_session_ptr(
 		new client_session(
 	    	shared_from_this(),
 			m_io_service,
-			m_def_con_handler
+			m_def_con_handler,
+			m_max_message_size*2
 		)
 	);
 	m_state = CLIENT_STATE_INITIALIZED;

src/websocket_client_session.cpp

@@ -47,8 +47,9 @@ using websocketpp::client_session;
 
 client_session::client_session (websocketpp::client_ptr c,
                                 boost::asio::io_service& io_service,
-                                websocketpp::connection_handler_ptr defc)
-	: session(io_service,defc),m_client(c) {}
+                                websocketpp::connection_handler_ptr defc,
+								uint64_t buf_size)
+	: session(io_service,defc,buf_size),m_client(c) {}
 
 void client_session::on_connect() {
 	// TODO: section 4.1: Figure out if we have another connection to this 
@@ -124,56 +125,49 @@ void client_session::read_handshake() {
 void client_session::handle_read_handshake(const boost::system::error_code& e,
 	std::size_t bytes_transferred) {
 	// parse server handshake
-	
-	
-	// read handshake and set local state (or pass to write_handshake)
-	std::ostringstream line;
-	line << &m_buf;
-	m_raw_server_handshake += line.str();
-	
-	m_buf << m_raw_server_handshake.substr(bytes_transferred);
-	
-	std::stringstream foo;
-	
-	foo << "data size: " << m_raw_server_handshake.size() 
-	<< " bytes transferred" << bytes_transferred;
-	
-	m_client->access_log(foo.str(), ALOG_HANDSHAKE);
-	m_client->access_log(m_raw_server_handshake,ALOG_HANDSHAKE);
-	m_client->access_log("SPACER",ALOG_HANDSHAKE);
-	
-	std::vector<std::string> tokens;
-	std::string::size_type start = 0;
+	std::istream response_stream(&m_buf);
+	std::string header;
 	std::string::size_type end;
 	
-	// Get request and parse headers
-	end = m_raw_server_handshake.find("\r\n",start);
-	
-	while(end != std::string::npos) {
-		tokens.push_back(m_raw_server_handshake.substr(start, end - start));
-		
-		start = end + 2;
-		
-		end = m_raw_server_handshake.find("\r\n",start);
+	// get status line
+	std::getline(response_stream, header);
+	if (header[header.size()-1] == '\r') {
+		header.erase(header.end()-1);
+		m_server_http_request = header;
+		m_raw_server_handshake += header+"\n";
 	}
 	
-	for (size_t i = 0; i < tokens.size(); i++) {
-		if (i == 0) {
-			m_server_http_request = tokens[i];
+	// get headers
+	while (std::getline(response_stream, header) && header != "\r") {
+		if (header[header.size()-1] != '\r') {
+			continue; // ignore malformed header lines?
+		} else {
+			header.erase(header.end()-1);
 		}
 		
-		end = tokens[i].find(": ",0);
+		end = header.find(": ",0);
 		
-		if (end != std::string::npos) {
-			std::string h = tokens[i].substr(0,end);
+		if (end != std::string::npos) {			
+			std::string h = header.substr(0,end);
 			if (get_server_header(h) == "") {
-				m_server_headers[h] = tokens[i].substr(end+2);
+				m_server_headers[h] = header.substr(end+2);
 			} else {
-				m_server_headers[h] += ", " + tokens[i].substr(end+2);
+				m_server_headers[h] += ", " + header.substr(end+2);
 			}
 		}
+		
+		m_raw_server_handshake += header+"\n";
 	}
 	
+	// temporary debugging
+	if (m_buf.size() > 0) {
+		std::stringstream foo;
+		foo << "bytes left over: " << m_buf.size();
+		access_log(foo.str(), ALOG_HANDSHAKE);
+	}
+	
+	m_client->access_log(m_raw_server_handshake,ALOG_HANDSHAKE);
+	
 	// handshake error checking
 	try {
 		std::stringstream err;
@@ -202,7 +196,7 @@ void client_session::handle_read_handshake(const boost::system::error_code& e,
 		if (h == "") {
 			throw(handshake_error("Required Upgrade header is missing",400));
 		} else if (!boost::iequals(h,"websocket")) {
-			err << "Upgrade header was " << h << " instead of \"websocket\"";
+			err << "Upgrade header was \"" << h << "\" instead of \"websocket\"";
 			throw(handshake_error(err.str(),400));
 		}
 		

src/websocket_client_session.hpp

@@ -65,7 +65,8 @@ class client_session : public session {
 public:
 	client_session (client_ptr c,
 			 boost::asio::io_service& io_service,
-			 connection_handler_ptr defc);
+			 connection_handler_ptr defc,
+			 uint64_t buf_size);
 
 	/*** CLIENT INTERFACE ***/
 	

src/websocket_server.cpp

@@ -138,9 +138,11 @@ void server::access_log(std::string msg,uint16_t level) {
 }
 
 void server::start_accept() {
+	// TODO: sanity check whether the session buffer size bound could be reduced
 	server_session_ptr new_session(new server_session(shared_from_this(),
 	                                                  m_io_service,
-	                                                  m_def_con_handler));
+	                                                  m_def_con_handler,
+													  m_max_message_size*2));
 	
 	m_acceptor.async_accept(
 		new_session->socket(),

src/websocket_server_session.cpp

@@ -45,8 +45,9 @@ using websocketpp::server_session;
 
 server_session::server_session(websocketpp::server_ptr s,
                                boost::asio::io_service& io_service,
-                               websocketpp::connection_handler_ptr defc)
-	: session(io_service,defc),m_server(s) {}
+                               websocketpp::connection_handler_ptr defc,
+							   uint64_t buf_size)
+	: session(io_service,defc,buf_size),m_server(s) {}
 
 void server_session::on_connect() {
 	read_handshake();
@@ -91,6 +92,7 @@ void server_session::select_extension(const std::string& val) {
 }
 
 void server_session::read_handshake() {
+	// TODO: pass maximum size to m_buf cosntructor to prevent DoS here
 	boost::asio::async_read_until(
 		m_socket,
 		m_buf,

src/websocket_server_session.hpp

@@ -65,7 +65,8 @@ class server_session : public session {
 public:
 	server_session (server_ptr s,
 			 boost::asio::io_service& io_service,
-			 connection_handler_ptr defc);
+			 connection_handler_ptr defc,
+		     uint64_t buf_size);
 
 	/*** SERVER INTERFACE ***/
 	

src/websocket_session.cpp

@@ -44,7 +44,8 @@
 using websocketpp::session;
 
 session::session (boost::asio::io_service& io_service,
-				  websocketpp::connection_handler_ptr defc)
+				  websocketpp::connection_handler_ptr defc,
+				  uint64_t buf_size)
 	: m_status(CONNECTING),
 	  m_local_close_code(CLOSE_STATUS_NO_STATUS),
 	  m_remote_close_code(CLOSE_STATUS_NO_STATUS),
@@ -54,8 +55,7 @@ session::session (boost::asio::io_service& io_service,
 	  m_socket(io_service),
 	  m_io_service(io_service),
 	  m_local_interface(defc),
-	  
-	  
+	  m_buf(buf_size), // maximum buffered (unconsumed) bytes from network
 	  m_utf8_state(utf8_validator::UTF8_ACCEPT),
 	  m_utf8_codepoint(0) {}
 

src/websocket_session.hpp

@@ -94,7 +94,8 @@ public:
 	static const uint16_t CLOSE_STATUS_EXTENSION_REQUIRE = 1010;
 
 	session (boost::asio::io_service& io_service,
-			 connection_handler_ptr defc);
+			 connection_handler_ptr defc,
+			 uint64_t buf_size);
 	
 	tcp::socket& socket();
 	boost::asio::io_service& io_service();