Update the SOCKS proxy to use the new RelayManager

Gemfile.lock

@@ -497,7 +497,7 @@ GEM
       rex-core
       rex-struct2
       rex-text
-    rex-core (0.1.35)
+    rex-core (0.1.36)
     rex-encoder (0.1.8)
       metasm
       rex-arch

lib/rex/proto/proxy/socks5/server.rb

@@ -1,6 +1,7 @@
 # -*- coding: binary -*-
 
 require 'thread'
+require 'rex/io/relay_manager'
 require 'rex/socket'
 
 module Rex
@@ -22,6 +23,7 @@ def initialize(opts={})
       @clients       = ::Array.new
       @running       = false
       @server_thread = nil
+      @relay_manager = Rex::IO::RelayManager.new
     end
 
     #
@@ -100,6 +102,7 @@ def remove_client(client)
     end
 
     attr_reader :opts
+    attr_reader :relay_manager
   end
 end
 end

lib/rex/proto/proxy/socks5/server_client.rb

@@ -12,61 +12,6 @@ module Proxy
 # A client connected to the proxy server.
 #
 module Socks5
-  #
-  # A mixin for a socket to perform a relay to another socket.
-  #
-  module TcpRelay
-    #
-    # TcpRelay data coming in from relay_sock to this socket.
-    #
-    def relay(relay_client, relay_sock)
-      @relay_client = relay_client
-      @relay_sock   = relay_sock
-      # start the relay thread (modified from Rex::IO::StreamAbstraction)
-      @relay_thread = Rex::ThreadFactory.spawn("SOCKS5ProxyServerTcpRelay", false) do
-        loop do
-          closed = false
-          buf    = nil
-
-          begin
-            s = Rex::ThreadSafe.select([@relay_sock], nil, nil, 0.2)
-            next if s.nil? || s[0].nil?
-          rescue
-            closed = true
-          end
-
-          unless closed
-            begin
-              buf = @relay_sock.sysread( 32768 )
-              closed = buf.nil?
-            rescue
-              closed = true
-            end
-          end
-
-          unless closed
-            total_sent   = 0
-            total_length = buf.length
-            while total_sent < total_length
-              begin
-                data = buf[total_sent, buf.length]
-                sent = self.write(data)
-                total_sent += sent if sent > 0
-              rescue
-                closed = true
-                break
-              end
-            end
-          end
-
-          if closed
-            @relay_client.stop
-            ::Thread.exit
-          end
-        end
-      end
-    end
-  end
 
   #
   # A client connected to the SOCKS5 server.
@@ -122,13 +67,13 @@ def start
 
           # handle the request
           handle_command
-        rescue => exception
+        rescue StandardError => e
           # respond with a general failure to the client
           response         = ResponsePacket.new
           response.command = REPLY_GENERAL_FAILURE
           @lsock.put(response.to_binary_s)
 
-          wlog("Client.start - #{$!}")
+          elog('ServerClient#start - encountered a problem while processing the client connection', error:e)
           self.stop
         end
       end
@@ -259,15 +204,11 @@ def handle_command_udp_associate(request)
     end
 
     #
-    # Setup the TcpRelay between lsock and rsock.
+    # Setup the relay between lsock and rsock.
     #
     def setup_tcp_relay
-      # setup the two way relay for full duplex io
-      @lsock.extend(TcpRelay)
-      @rsock.extend(TcpRelay)
-      # start the socket relays...
-      @lsock.relay(self, @rsock)
-      @rsock.relay(self, @lsock)
+      @server.relay_manager.add_relay(@rsock, sink: @lsock, name: 'SOCKS5ProxyRelay-Remote', on_exit: method(:stop))
+      @server.relay_manager.add_relay(@lsock, sink: @rsock, name: 'SOCKS5ProxyRelay-Local', on_exit: method(:stop))
     end
 
     #
@@ -286,9 +227,13 @@ def stop
           rescue
           end
 
-          @client_thread.kill if @client_thread and @client_thread.alive?
           @server.remove_client(self)
           @closed = true
+
+          unless @client_thread == Thread.current
+            @client_thread.join
+          end
+          @client_thread = nil
         end
       end
     end

modules/payloads/singles/windows/x64/encrypted_shell_reverse_tcp.rb

@@ -4,7 +4,7 @@
 ##
 
 module MetasploitModule
-  CachedSize = 4176
+  CachedSize = 4608
 
   include Msf::Payload::Windows
   include Msf::Payload::Single

modules/payloads/stagers/windows/x64/encrypted_reverse_tcp.rb

@@ -4,7 +4,7 @@
 ##
 
 module MetasploitModule
-  CachedSize = 2688
+  CachedSize = 3072
 
   include Msf::Payload::Stager
   include Msf::Payload::Windows::EncryptedReverseTcp

spec/lib/rex/proto/proxy/socks5/server_client_spec.rb

@@ -0,0 +1,134 @@
+# -*- coding:binary -*-
+
+RSpec.describe Rex::Proto::Proxy::Socks5::ServerClient do
+  let(:relay_manager) { instance_double(Rex::IO::RelayManager) }
+  let(:server) do
+    instance_double(
+      Rex::Proto::Proxy::Socks5::Server,
+      opts: {},
+      add_client: nil,
+      remove_client: nil,
+      relay_manager: relay_manager
+    )
+  end
+  let(:lsock) { instance_double('Rex::Socket::Tcp') }
+  let(:rsock) { instance_double('Rex::Socket::Tcp') }
+
+  subject(:client) { described_class.new(server, lsock) }
+
+  before do
+    allow(lsock).to receive(:close)
+    allow(rsock).to receive(:close)
+    client.instance_variable_set(:@rsock, rsock)
+  end
+
+  describe '#setup_tcp_relay' do
+    before do
+      allow(relay_manager).to receive(:add_relay)
+    end
+
+    it 'registers a relay from rsock to lsock' do
+      expect(relay_manager).to receive(:add_relay).with(
+        rsock,
+        sink: lsock,
+        name: 'SOCKS5ProxyRelay-Remote',
+        on_exit: anything
+      )
+      client.send(:setup_tcp_relay)
+    end
+
+    it 'registers a relay from lsock to rsock' do
+      expect(relay_manager).to receive(:add_relay).with(
+        lsock,
+        sink: rsock,
+        name: 'SOCKS5ProxyRelay-Local',
+        on_exit: anything
+      )
+      client.send(:setup_tcp_relay)
+    end
+
+    it 'passes a callable stop callback for each relay' do
+      on_exit_callbacks = []
+      allow(relay_manager).to receive(:add_relay) do |_sock, **kwargs|
+        on_exit_callbacks << kwargs[:on_exit]
+      end
+
+      client.send(:setup_tcp_relay)
+
+      expect(on_exit_callbacks.size).to eq(2)
+      on_exit_callbacks.each { |cb| expect(cb).to respond_to(:call) }
+    end
+
+    it 'passes the client #stop method as the on_exit callback' do
+      on_exit_callbacks = []
+      allow(relay_manager).to receive(:add_relay) do |_sock, **kwargs|
+        on_exit_callbacks << kwargs[:on_exit]
+      end
+
+      client.send(:setup_tcp_relay)
+
+      on_exit_callbacks.each do |cb|
+        expect(cb).to eq(client.method(:stop))
+      end
+    end
+  end
+
+  describe '#stop' do
+    # Use a thread double so join doesn't actually block.
+    let(:mock_thread) do
+      thread = instance_double(Thread)
+      allow(thread).to receive(:join)
+      thread
+    end
+
+    before do
+      client.instance_variable_set(:@client_thread, mock_thread)
+    end
+
+    it 'closes the local socket' do
+      expect(lsock).to receive(:close)
+      client.stop
+    end
+
+    it 'closes the remote socket' do
+      expect(rsock).to receive(:close)
+      client.stop
+    end
+
+    it 'removes itself from the server' do
+      expect(server).to receive(:remove_client).with(client)
+      client.stop
+    end
+
+    it 'joins the client thread when called from an external thread' do
+      expect(mock_thread).to receive(:join)
+      client.stop
+    end
+
+    it 'does not deadlock when called from within the client thread' do
+      completed = false
+
+      # Use a real thread and have it set @client_thread to itself before calling stop
+      client_thread = Thread.new do
+        client.instance_variable_set(:@client_thread, Thread.current)
+        client.stop
+        completed = true
+      end
+
+      result = client_thread.join(2)
+      expect(result).not_to be_nil, 'stop deadlocked when called from the client thread'
+      expect(completed).to be(true)
+    end
+
+    it 'is idempotent - a second call is a no-op' do
+      client.stop
+      expect { client.stop }.not_to raise_error
+      expect(lsock).to have_received(:close).once
+    end
+
+    it 'sets @closed so subsequent calls are skipped' do
+      client.stop
+      expect(client.instance_variable_get(:@closed)).to be(true)
+    end
+  end
+end

spec/lib/rex/proto/proxy/socks5/server_spec.rb

@@ -6,11 +6,75 @@
     Rex::Proto::Proxy::Socks5::Server.new
   end
 
+  describe "#initialize" do
+    it "defaults ServerHost to 0.0.0.0" do
+      expect(server.opts['ServerHost']).to eq('0.0.0.0')
+    end
+
+    it "defaults ServerPort to 1080" do
+      expect(server.opts['ServerPort']).to eq(1080)
+    end
+
+    it "accepts option overrides" do
+      s = described_class.new('ServerPort' => 9090)
+      expect(s.opts['ServerPort']).to eq(9090)
+    end
+
+    it "starts in a not-running state" do
+      expect(server.is_running?).to be(false)
+    end
+  end
+
   describe "#is_running?" do
 
     it "should respond to #is_running?" do
       expect(server.is_running?).to eq(false)
     end
 
   end
+
+  describe "#relay_manager" do
+    it "returns a Rex::IO::RelayManager instance" do
+      expect(server.relay_manager).to be_a(Rex::IO::RelayManager)
+    end
+
+    it "returns the same instance on repeated calls" do
+      expect(server.relay_manager).to be(server.relay_manager)
+    end
+  end
+
+  describe "#add_client" do
+    let(:client) { double('ServerClient') }
+
+    it "starts with no clients" do
+      expect(server.instance_variable_get(:@clients)).to be_empty
+    end
+
+    it "adds a client" do
+      server.add_client(client)
+      expect(server.instance_variable_get(:@clients)).to include(client)
+    end
+  end
+
+  describe "#remove_client" do
+    let(:client) { double('ServerClient') }
+
+    it "removes a client" do
+      server.add_client(client)
+      server.remove_client(client)
+      expect(server.instance_variable_get(:@clients)).not_to include(client)
+    end
+  end
+
+  describe "#stop" do
+    context "when the server is not running" do
+      it "returns true" do
+        expect(server.stop).to be(true)
+      end
+
+      it "does not raise" do
+        expect { server.stop }.not_to raise_error
+      end
+    end
+  end
 end