From b60990c19cc8646a754d7aa07b61e69991d28cc6 Mon Sep 17 00:00:00 2001
From: RageLtMan <rageltman [at] sempervictus>
Date: Mon, 21 Mar 2016 13:53:44 -0400
Subject: [PATCH] Use a MockDnsClient object for request state

In order to handle TCP and UDP clients in a common manner, the
DNS server created a Rex::Socket::Udp object to represent the
client object allowing for a client.write(response) approach to
returning results for both TCP and UDP clients. During work on
the common socket abstractions (#6692) it became apparent that
remote pivoted sockets cannot be created with the same exact param
set used on the server socket - sockets dont reuse with localhost
and localport params being the same, an exception is raised from
the Windows side of the pivot abstraction. Creating a new socket
for every request is also needless overhead and noise.

Create the MockDnsClient class to  consume peerhost, peerport, and
the DNS server's UDP socket as arguments in order to execute a
sendto() from the existing socket when sending a response. A write
method is provided in the class for common interface between the
UDP and TCP request handlers.

This has been tested in conjunction with #6692 and shown to be
successful as serving remote requests from the IO.select polled
pivot socket running on a Windows host via Meterpreter.
---
 lib/rex/proto/dns/server.rb | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/lib/rex/proto/dns/server.rb b/lib/rex/proto/dns/server.rb
index 1f10e9d12f..2289f692c4 100644
--- a/lib/rex/proto/dns/server.rb
+++ b/lib/rex/proto/dns/server.rb
@@ -127,6 +127,32 @@ class Server
     end
   end # Cache
 
+  class MockDnsClient
+    attr_reader :peerhost, :peerport, :srvsock
+
+    #
+    # Create mock DNS client
+    #
+    # @param host [String] PeerHost IP address
+    # @param port [Fixnum] PeerPort integer
+    def initialize(host, port, sock)
+      @peerhost = host
+      @peerport = port
+      @srvsock = sock
+    end
+
+    #
+    # Test method to prevent GC/ObjectSpace abuse via class lookups
+    #
+    def mock_dns_client?
+      true
+    end
+
+    def write(data)
+      srvsock.sendto(data, peerhost, peerport)
+    end
+  end
+
   include Rex::IO::GramServer
 
   Packet = Rex::Proto::DNS::Packet
@@ -292,13 +318,8 @@ protected
 
       if (r != nil and r[0] == self.udp_sock)
         buf,host,port = self.udp_sock.recvfrom(65535)
-        # Mock up a client object as a Rex Socket for sending back data
-        cli = Rex::Socket::Udp.create(
-          'PeerHost' => host,
-          'PeerPort' => port,
-          'LocalHost' => self.udp_sock.localhost,
-          'LocalPort' => self.udp_sock.localport
-        )
+        # Mock up a client object for sending back data
+        cli = MockDnsClient.new(host, port, r[0])
         dispatch_request(cli, buf)
       end
     end