Improve validation of Websocket requests

lib/async/websocket/connect_request.rb

@@ -12,7 +12,7 @@
 
 module Async
 	module WebSocket
-		# This is required for HTTP/1.x to upgrade the connection to the WebSocket protocol.
+		# This is required for HTTP/2 to establish a connection using the WebSocket protocol.
 		# See https://tools.ietf.org/html/rfc8441 for more details.
 		class ConnectRequest < ::Protocol::HTTP::Request
 			include ::Protocol::WebSocket::Headers

lib/async/websocket/upgrade_request.rb

@@ -17,6 +17,7 @@
 module Async
 	module WebSocket
 		# This is required for HTTP/1.x to upgrade the connection to the WebSocket protocol.
+		# See https://tools.ietf.org/html/rfc6455
 		class UpgradeRequest < ::Protocol::HTTP::Request
 			include ::Protocol::WebSocket::Headers
 
@@ -75,8 +76,9 @@ def call(connection)
 						raise ProtocolError, "Invalid accept digest, expected #{expected_accept_digest.inspect}, got #{accept_digest.inspect}!"
 					end
 				end
+				verified = accept_digest && Array(response.protocol) == %w(websocket) && response.headers['connection']&.include?('upgrade')
 
-				return Wrapper.new(response, verified: !!accept_digest)
+				return Wrapper.new(response, verified: verified)
 			end
 		end
 	end

test/async/websocket/client.rb

@@ -105,17 +105,53 @@
 	end
 end
 
+FailedToNegotiate = Sus::Shared("a failed websocket request") do
+	it 'raises an error' do
+		expect do
+			Async::WebSocket::Client.connect(client_endpoint) {}
+		end.to raise_exception(Async::WebSocket::ProtocolError, message: be =~ /Failed to negotiate connection/)
+	end
+end
+
 describe Async::WebSocket::Client do
 	include Sus::Fixtures::Async::HTTP::ServerContext
 
 	with 'http/1' do
 		let(:protocol) {Async::HTTP::Protocol::HTTP1}
 		it_behaves_like ClientExamples
 
+		def valid_headers(request)
+			{
+				'connection' => 'upgrade',
+				'upgrade' => 'websocket',
+				'sec-websocket-accept' => Protocol::WebSocket::Headers::Nounce.accept_digest(request.headers['sec-websocket-key'].first)
+			}
+		end
+
+		with 'invalid connection header' do
+			let(:app) do
+				Protocol::HTTP::Middleware.for do |request|
+					Protocol::HTTP::Response[101, valid_headers(request).except('connection'), []]
+				end
+			end
+
+			it_behaves_like FailedToNegotiate
+		end
+
+		with 'invalid upgrade header' do
+			let(:app) do
+				Protocol::HTTP::Middleware.for do |request|
+					Protocol::HTTP::Response[101, valid_headers(request).except('upgrade'), []]
+				end
+			end
+
+			it_behaves_like FailedToNegotiate
+		end
+
 		with 'invalid sec-websocket-accept header' do
 			let(:app) do
 				Protocol::HTTP::Middleware.for do |request|
-					Protocol::HTTP::Response[101, {'sec-websocket-accept'=>'wrong-digest'}, []]
+					Protocol::HTTP::Response[101, valid_headers(request).merge('sec-websocket-accept'=>'wrong-digest'), []]
 				end
 			end
 
@@ -125,24 +161,40 @@
 				end.to raise_exception(Async::WebSocket::ProtocolError, message: be =~ /Invalid accept digest/)
 			end
 		end
-			
+
 		with 'missing sec-websocket-accept header' do
 			let(:app) do
 				Protocol::HTTP::Middleware.for do |request|
-					Protocol::HTTP::Response[101, {}, []]
+					Protocol::HTTP::Response[101, valid_headers(request).except('sec-websocket-accept'), []]
 				end
 			end
 
-			it 'raises an error' do
-				expect do
-					Async::WebSocket::Client.connect(client_endpoint) {}
-				end.to raise_exception(Async::WebSocket::ProtocolError, message: be =~ /Failed to negotiate connection/)
+			it_behaves_like FailedToNegotiate
+		end
+
+		with 'invalid status' do
+			let(:app) do
+				Protocol::HTTP::Middleware.for do |request|
+					Protocol::HTTP::Response[403, valid_headers(request), []]
+				end
 			end
+
+			it_behaves_like FailedToNegotiate
 		end
 	end
 
 	with 'http/2' do
 		let(:protocol) {Async::HTTP::Protocol::HTTP2}
 		it_behaves_like ClientExamples
+
+		with 'invalid status' do
+			let(:app) do
+				Protocol::HTTP::Middleware.for do |request|
+					Protocol::HTTP::Response[403, {}, []]
+				end
+			end
+
+			it_behaves_like FailedToNegotiate
+		end
 	end
 end