From 6146dd1d025640867abd5503240d181ca325adb8 Mon Sep 17 00:00:00 2001
From: zarqman <tm@iprog.com>
Date: Sat, 9 Dec 2023 15:05:43 -0700
Subject: [PATCH] Improve validation of Websocket requests. (#59)

* Verify `connection` and `upgrade` headers for HTTP/1 requests.

* Add test for status for HTTP/2.
---
 lib/async/websocket/connect_request.rb |  2 +-
 lib/async/websocket/upgrade_request.rb |  4 +-
 test/async/websocket/client.rb         | 66 +++++++++++++++++++++++---
 3 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/lib/async/websocket/connect_request.rb b/lib/async/websocket/connect_request.rb
index 05e00f8..7495d23 100644
--- a/lib/async/websocket/connect_request.rb
+++ b/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
diff --git a/lib/async/websocket/upgrade_request.rb b/lib/async/websocket/upgrade_request.rb
index cadbebb..817a37b 100644
--- a/lib/async/websocket/upgrade_request.rb
+++ b/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
diff --git a/test/async/websocket/client.rb b/test/async/websocket/client.rb
index 64d7bb0..663be6b 100644
--- a/test/async/websocket/client.rb
+++ b/test/async/websocket/client.rb
@@ -105,6 +105,14 @@
 	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
 
@@ -112,10 +120,38 @@
 		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