ssl_socket.rb in async-io-1.5.0

- old
+ new
@@ -26,72 +26,104 @@
 	module IO
 		SSLError = OpenSSL::SSL::SSLError
 		
 		# Asynchronous TCP socket wrapper.
 		class SSLSocket < Generic
-			wraps ::OpenSSL::SSL::SSLSocket, :alpn_protocol, :cert, :cipher, :client_ca, :hostname=, :npn_protocol, :peer_cert, :peer_cert_chain, :pending, :post_connection_check, :session, :session=, :session_reused?, :ssl_version, :state
+			wraps ::OpenSSL::SSL::SSLSocket, :alpn_protocol, :cert, :cipher, :client_ca, :close, :context, :hostname, :hostname=, :npn_protocol, :peer_cert, :peer_cert_chain, :pending, :post_connection_check, :session, :session=, :session_reused?, :ssl_version, :state, :sync_close, :sync_close=, :sysclose, :verify_result, :tmp_key
 			
 			wrap_blocking_method :accept, :accept_nonblock
 			wrap_blocking_method :connect, :connect_nonblock
 			
+			alias syswrite write
+			alias sysread read
+			
+			# It's hard to know what #to_io / #io should do. So, they are omitted.
+			
+			def self.connect(socket, context, hostname = nil, &block)
+				client = self.wrap(socket, context)
+				
+				# Used for SNI:
+				if hostname
+					client.hostname = hostname
+				end
+				
+				begin
+					client.connect
+				rescue
+					# If the connection fails (e.g. certificates are invalid), the caller never sees the socket, so we close it and raise the exception up the chain.
+					client.close
+					
+					raise
+				end
+				
+				return client unless block_given?
+				
+				begin
+					yield client
+				ensure
+					client.close
+				end
+			end
+			
 			def local_address
 				@io.to_io.local_address
 			end
 			
 			def remote_address
 				@io.to_io.remote_address
 			end
 			
-			# This method/implementation might change in the future, don't depend on it :)
-			def self.connect_socket(socket, context)
-				io = wrapped_klass.new(socket.to_io, context)
+			def self.wrap(socket, context)
+				io = @wrapped_klass.new(socket.to_io, context)
 				
 				# This ensures that when the internal IO is closed, it also closes the internal socket:
 				io.sync_close = true
 				
 				return self.new(io, socket.reactor)
 			end
 		end
 		
+		# We reimplement this from scratch because the native implementation doesn't expose the underlying server/context that we need to implement non-blocking accept.
 		class SSLServer
 			extend Forwardable
 			
 			def initialize(server, context)
 				@server = server
 				@context = context
 			end
 			
-			def_delegators :@server, :local_address, :setsockopt, :getsockopt
+			def_delegators :@server, :local_address, :setsockopt, :getsockopt, :close
 			
 			attr :server
 			attr :context
 			
-			include ServerSocket
-			
 			def listen(*args)
 				@server.listen(*args)
 			end
 			
 			def accept(task: Task.current)
 				peer, address = @server.accept
 				
-				wrapper = SSLSocket.connect_socket(peer, @context)
+				wrapper = SSLSocket.wrap(peer, @context)
 				
 				return wrapper, address unless block_given?
 				
 				task.async do
 					task.annotate "accepting secure connection #{address}"
 					
 					begin
+						# You want to do this in a nested async task or you might suffer from head-of-line blocking.
 						wrapper.accept
 						
 						yield wrapper, address
 					rescue SSLError
 						Async.logger.error($!.class) {$!}
 					ensure
 						wrapper.close
 					end
 				end
 			end
+			
+			include Server
 		end
 	end
 end