Add the #starttls method for TCP sockets

lib/rex/socket.rb

@@ -760,6 +760,9 @@ def self.tcp_socket_pair
       raise "Thread #{i} - error #{e} - last child error: #{last_child_error}"
     end
 
+    lsock.extend(Rex::Socket::Tcp)
+    rsock.extend(Rex::Socket::Tcp)
+
     return [lsock, rsock]
   end
 
@@ -779,6 +782,9 @@ def self.udp_socket_pair
 
     lsock.connect( *rsock.addr.values_at(3,1) )
 
+    lsock.extend(Rex::Socket::Udp)
+    rsock.extend(Rex::Socket::Udp)
+
     return [lsock, rsock]
   end
 

lib/rex/socket/comm/local.rb

@@ -324,8 +324,7 @@ def self.create_by_type(param, type, proto = 0)
 
       # Now extend the socket with SSL and perform the handshake
       if !param.bare? && param.ssl
-        klass = Rex::Socket::SslTcp
-        sock.extend(klass)
+        sock.extend(Rex::Socket::SslTcp)
         sock.initsock(param)
       end
     end

lib/rex/socket/ssl_tcp.rb

@@ -66,21 +66,19 @@ def initsock(params = nil)
     super
 
     version = params&.ssl_version || Rex::Socket::Ssl::DEFAULT_SSL_VERSION
-    # Raise an error if no selected versions are supported
-    unless Rex::Socket::SslTcp.system_ssl_methods.include? version
-      raise ArgumentError,
-        "This version of Ruby does not support the requested SSL/TLS version #{version}"
-    end
 
     # Try initializing the socket with this SSL/TLS version
     # This will throw an exception if it fails
     initsock_with_ssl_version(params, version)
-
-    # Track the SSL version
-    self.ssl_negotiated_version = version
   end
 
   def initsock_with_ssl_version(params, version)
+    # Raise an error if no selected versions are supported
+    unless Rex::Socket::SslTcp.system_ssl_methods.include? version
+      raise ArgumentError,
+        "This version of Ruby does not support the requested SSL/TLS version #{version}"
+    end
+
     # Build the SSL connection
     self.sslctx  = OpenSSL::SSL::SSLContext.new(version)
 
@@ -149,33 +147,22 @@ def initsock_with_ssl_version(params, version)
 
     # Force a negotiation timeout
     begin
-    Timeout.timeout(params.timeout) do
-      if not allow_nonblock?
-        self.sslsock.connect
-      else
-        begin
-          self.sslsock.connect_nonblock
-        # Ruby 1.8.7 and 1.9.0/1.9.1 uses a standard Errno
-        rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
-            IO::select(nil, nil, nil, 0.10)
-            retry
-
-        # Ruby 1.9.2+ uses IO::WaitReadable/IO::WaitWritable
-        rescue ::Exception => e
-          if ::IO.const_defined?('WaitReadable') and e.kind_of?(::IO::WaitReadable)
+      Timeout.timeout(params.timeout) do
+        if not allow_nonblock?
+          self.sslsock.connect
+        else
+          begin
+            self.sslsock.connect_nonblock
+          rescue ::IO::WaitReadable
             IO::select( [ self.sslsock ], nil, nil, 0.10 )
             retry
-          end
 
-          if ::IO.const_defined?('WaitWritable') and e.kind_of?(::IO::WaitWritable)
+          rescue ::IO::WaitWritable
             IO::select( nil, [ self.sslsock ], nil, 0.10 )
             retry
           end
-
-          raise e
         end
       end
-    end
 
     rescue ::Timeout::Error
       raise Rex::ConnectionTimeout.new(params.peerhost, params.peerport)
@@ -215,34 +202,16 @@ def write(buf, opts = {})
     rescue ::IOError, ::Errno::EPIPE
       return nil
 
-    # Ruby 1.8.7 and 1.9.0/1.9.1 uses a standard Errno
-    rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
-      # Sleep for a half a second, or until we can write again
-      Rex::ThreadSafe.select( nil, [ self.sslsock ], nil, retry_time )
-      # Decrement the block size to handle full sendQs better
-      block_size = 1024
-      # Try to write the data again
+    rescue ::IO::WaitReadable
+      IO::select( [ self.sslsock ], nil, nil, retry_time )
       retry
 
-    # Ruby 1.9.2+ uses IO::WaitReadable/IO::WaitWritable
-    rescue ::Exception => e
-      if ::IO.const_defined?('WaitReadable') and e.kind_of?(::IO::WaitReadable)
-        IO::select( [ self.sslsock ], nil, nil, retry_time )
-        retry
-      end
-
-      if ::IO.const_defined?('WaitWritable') and e.kind_of?(::IO::WaitWritable)
-        IO::select( nil, [ self.sslsock ], nil, retry_time )
-        retry
-      end
-
-      # Another form of SSL error, this is always fatal
-      if e.kind_of?(::OpenSSL::SSL::SSLError)
-        return nil
-      end
+    rescue ::IO::WaitWritable
+      IO::select( nil, [ self.sslsock ], nil, retry_time )
+      retry
 
-      # Bubble the event up to the caller otherwise
-      raise e
+    rescue ::OpenSSL::SSL::SSLError
+      return nil
     end
 
     total_sent
@@ -298,33 +267,16 @@ def read(length = nil, opts = {})
     rescue ::IOError, ::Errno::EPIPE
       return nil
 
-    # Ruby 1.8.7 and 1.9.0/1.9.1 uses a standard Errno
-    rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
-      # Sleep for a tenth a second, or until we can read again
-      Rex::ThreadSafe.select( [ self.sslsock ], nil, nil, 0.10 )
-      # Decrement the block size to handle full sendQs better
-      block_size = 1024
-      # Try to write the data again
+    rescue ::IO::WaitReadable
+      IO::select( [ self.sslsock ], nil, nil, 0.10 )
       retry
 
-    # Ruby 1.9.2+ uses IO::WaitReadable/IO::WaitWritable
-    rescue ::Exception => e
-      if ::IO.const_defined?('WaitReadable') and e.kind_of?(::IO::WaitReadable)
-        IO::select( [ self.sslsock ], nil, nil, 0.5 )
-        retry
-      end
-
-      if ::IO.const_defined?('WaitWritable') and e.kind_of?(::IO::WaitWritable)
-        IO::select( nil, [ self.sslsock ], nil, 0.5 )
-        retry
-      end
-
-      # Another form of SSL error, this is always fatal
-      if e.kind_of?(::OpenSSL::SSL::SSLError)
-        return nil
-      end
+    rescue ::IO::WaitWritable
+      IO::select( nil, [ self.sslsock ], nil, 0.10 )
+      retry
 
-      raise e
+    rescue ::OpenSSL::SSL::SSLError
+      return nil
     end
 
   end
@@ -409,7 +361,6 @@ def allow_nonblock?
   end
 
   attr_reader :peer_verified # :nodoc:
-  attr_reader :ssl_negotiated_version # :nodoc:
   attr_accessor :sslsock, :sslctx, :sslhash # :nodoc:
 
   def type?
@@ -419,8 +370,6 @@ def type?
 protected
 
   attr_writer :peer_verified # :nodoc:
-  attr_writer :ssl_negotiated_version # :nodoc:
-
 
 rescue LoadError
 end

lib/rex/socket/ssl_tcp_server.rb

@@ -71,24 +71,13 @@ def accept(opts = {})
         begin
           ssl.accept_nonblock
 
-        # Ruby 1.8.7 and 1.9.0/1.9.1 uses a standard Errno
-        rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
-            IO::select(nil, nil, nil, 0.10)
-            retry
-
-        # Ruby 1.9.2+ uses IO::WaitReadable/IO::WaitWritable
-        rescue ::Exception => e
-          if ::IO.const_defined?('WaitReadable') and e.kind_of?(::IO::WaitReadable)
-            IO::select( [ ssl ], nil, nil, 0.10 )
-            retry
-          end
-
-          if ::IO.const_defined?('WaitWritable') and e.kind_of?(::IO::WaitWritable)
-            IO::select( nil, [ ssl ], nil, 0.10 )
-            retry
-          end
-
-          raise e
+        rescue ::IO::WaitReadable
+          IO::select( [ self.sslsock ], nil, nil, 0.10 )
+          retry
+
+        rescue ::IO::WaitWritable
+          IO::select( nil, [ self.sslsock ], nil, 0.10 )
+          retry
         end
       end
 

lib/rex/socket/tcp.rb

@@ -58,4 +58,12 @@ def type?
     return 'tcp'
   end
 
+  def starttls(param)
+    param = Rex::Socket::Parameters.from_hash(param) if param.is_a? Hash
+
+    param.ssl = true
+    extend(Rex::Socket::SslTcp)
+    initsock_with_ssl_version(param, (param.ssl_version || Rex::Socket::Ssl::DEFAULT_SSL_VERSION))
+    nil
+  end
 end

spec/rex/socket/tcp_spec.rb

@@ -0,0 +1,21 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+RSpec.describe Rex::Socket::Tcp do
+  describe '#starttls' do
+    it 'calls Parameters.to_hash with a hash argument' do
+      socket = described_class.create
+      expect(Rex::Socket::Parameters).to receive(:from_hash).with({'SSL' => true}).and_return(Rex::Socket::Parameters.new)
+      expect(socket).to receive(:initsock_with_ssl_version).and_return(nil)
+      socket.starttls('SSL' => true)
+    end
+
+    it 'accepts Parameters as an argument' do
+      socket = described_class.create
+      parameters = Rex::Socket::Parameters.new
+      expect(Rex::Socket::Parameters).to_not receive(:new)
+      expect(socket).to receive(:initsock_with_ssl_version).with(parameters, Rex::Socket::Ssl::DEFAULT_SSL_VERSION).and_return(nil)
+      socket.starttls(parameters)
+    end
+  end
+end

spec/rex/socket_spec.rb

@@ -28,8 +28,6 @@
 
     it 'creates two socket pairs' do
       lsock, rsock = described_class.tcp_socket_pair
-      lsock.extend(Rex::IO::Stream)
-      rsock.extend(Rex::IO::Stream)
 
       expect(lsock.closed?).to be(false)
       expect(rsock.closed?).to be(false)
@@ -40,6 +38,33 @@
       expect(lsock.closed?).to be(true)
       expect(rsock.closed?).to be(true)
     end
+
+    it 'extends returned with the Rex::Socket API' do
+      lsock, rsock = described_class.tcp_socket_pair
+      expect(lsock).to be_a(Rex::Socket::Tcp)
+      expect(rsock).to be_a(Rex::Socket::Tcp)
+    end
+  end
+
+  describe '.udp_socket_pair' do
+    it 'creates two socket pairs' do
+      lsock, rsock = described_class.udp_socket_pair
+
+      expect(lsock.closed?).to be(false)
+      expect(rsock.closed?).to be(false)
+
+      lsock.close
+      rsock.close
+
+      expect(lsock.closed?).to be(true)
+      expect(rsock.closed?).to be(true)
+    end
+
+    it 'extends returned with the Rex::Socket API' do
+      lsock, rsock = described_class.udp_socket_pair
+      expect(lsock).to be_a(Rex::Socket::Udp)
+      expect(rsock).to be_a(Rex::Socket::Udp)
+    end
   end
 
   describe '.addr_itoa' do