From e5d0370a949d940e75e0b9fda078d2bcf5589ffb Mon Sep 17 00:00:00 2001
From: Pearce Barry <pearce_barry@rapid7.com>
Date: Tue, 21 Feb 2017 15:05:20 -0600
Subject: [PATCH] Fixes MS-1716, keep sessions in progress alive.

---
 lib/rex/post/meterpreter/packet_dispatcher.rb | 46 +++++++++++++------
 lib/rex/post/meterpreter/packet_parser.rb     | 21 +++++----
 .../meterpreter/packet_response_waiter.rb     | 27 ++++++++---
 .../post/meterpreter/packet_parser_spec.rb    |  3 +-
 4 files changed, 69 insertions(+), 28 deletions(-)

diff --git a/lib/rex/post/meterpreter/packet_dispatcher.rb b/lib/rex/post/meterpreter/packet_dispatcher.rb
index cb4a95702b..21bd21ac35 100644
--- a/lib/rex/post/meterpreter/packet_dispatcher.rb
+++ b/lib/rex/post/meterpreter/packet_dispatcher.rb
@@ -284,6 +284,20 @@ module PacketDispatcher
   # Reception
   #
   ##
+
+  #
+  # Simple class to track packets and if they are in-progress or complete.
+  #
+  class QueuedPacket
+    attr_reader :packet
+    attr_reader :in_progress
+
+    def initialize(packet, in_progress)
+      @packet = packet
+      @in_progress = in_progress
+    end
+  end
+
   #
   # Monitors the PacketDispatcher's sock for data in its own
   # thread context and parsers all inbound packets.
@@ -306,8 +320,8 @@ module PacketDispatcher
         begin
           rv = Rex::ThreadSafe.select([ self.sock.fd ], nil, nil, PING_TIME)
           if rv
-            packet = receive_packet
-            @pqueue << packet if packet
+            packet, in_progress = receive_packet
+            @pqueue << QueuedPacket.new(packet, in_progress)
           elsif self.send_keepalives && @pqueue.empty?
             keepalive
           end
@@ -342,11 +356,11 @@ module PacketDispatcher
         tmp_channel = []
         tmp_close   = []
         backlog.each do |pkt|
-          if(pkt.response?)
+          if(pkt.packet.response?)
             tmp_command << pkt
             next
           end
-          if(pkt.method == "core_channel_close")
+          if(pkt.packet.method == "core_channel_close")
             tmp_close << pkt
             next
           end
@@ -365,7 +379,7 @@ module PacketDispatcher
         backlog.each do |pkt|
 
           begin
-          if ! dispatch_inbound_packet(pkt)
+          if ! dispatch_inbound_packet(pkt.packet, pkt.in_progress)
             # Keep Packets in the receive queue until a handler is registered
             # for them. Packets will live in the receive queue for up to
             # PACKET_TIMEOUT seconds, after which they will be dropped.
@@ -373,13 +387,15 @@ module PacketDispatcher
             # A common reason why there would not immediately be a handler for
             # a received Packet is in channels, where a connection may
             # open and receive data before anything has asked to read.
-            if (::Time.now.to_i - pkt.created_at.to_i < PACKET_TIMEOUT)
+            #
+            # Also, don't bother saving incomplete packets if we have no handler.
+            if (!pkt.in_progress and ::Time.now.to_i - pkt.packet.created_at.to_i < PACKET_TIMEOUT)
               incomplete << pkt
             end
           end
 
           rescue ::Exception => e
-            dlog("Dispatching exception with packet #{pkt}: #{e} #{e.backtrace}", 'meterpreter', LEV_1)
+            dlog("Dispatching exception with packet #{pkt.packet}: #{e} #{e.backtrace}", 'meterpreter', LEV_1)
           end
         end
 
@@ -459,12 +475,16 @@ module PacketDispatcher
   # Notifies a whomever is waiting for a the supplied response,
   # if anyone.
   #
-  def notify_response_waiter(response)
+  # For not-yet-complete responses, we might not be able to determine
+  # the response ID, in that case just let all waiters know that some
+  # responses are trickling in.
+  #
+  def notify_response_waiter(response, in_progress=false)
     handled = false
     self.waiters.each() { |waiter|
-      if (waiter.waiting_for?(response))
-        waiter.notify(response)
-        remove_response_waiter(waiter)
+      if (in_progress || waiter.waiting_for?(response))
+        waiter.notify(response, in_progress)
+        remove_response_waiter(waiter) unless in_progress
         handled = true
         break
       end
@@ -498,7 +518,7 @@ module PacketDispatcher
   # Otherwise, the packet is passed onto any registered dispatch
   # handlers until one returns success.
   #
-  def dispatch_inbound_packet(packet)
+  def dispatch_inbound_packet(packet, in_progress=false)
     handled = false
 
     # Update our last reply time
@@ -507,7 +527,7 @@ module PacketDispatcher
     # If the packet is a response, try to notify any potential
     # waiters
     if packet.response?
-      if (notify_response_waiter(packet))
+      if (notify_response_waiter(packet, in_progress))
         return true
       end
     end
diff --git a/lib/rex/post/meterpreter/packet_parser.rb b/lib/rex/post/meterpreter/packet_parser.rb
index 5b33c7b7c5..4575381790 100644
--- a/lib/rex/post/meterpreter/packet_parser.rb
+++ b/lib/rex/post/meterpreter/packet_parser.rb
@@ -75,22 +75,27 @@ class PacketParser
       end
     end
 
+    in_progress = true
+
+    # TODO: cipher decryption
+    if (cipher)
+    end
+
+    # Deserialize the packet from the raw buffer
+    packet.from_r(self.raw)
+
     # If we've finished reading the entire packet
     if ((self.hdr_length_left == 0) &&
         (self.payload_length_left == 0))
 
-      # TODO: cipher decryption
-      if (cipher)
-      end
-
-      # Deserialize the packet from the raw buffer
-      packet.from_r(self.raw)
-
       # Reset our state
       reset
 
-      return packet
+      # packet is complete!
+      in_progress = false
     end
+
+    return packet, in_progress
   end
 
 protected
diff --git a/lib/rex/post/meterpreter/packet_response_waiter.rb b/lib/rex/post/meterpreter/packet_response_waiter.rb
index 5f2f1557d7..637e1b9a95 100644
--- a/lib/rex/post/meterpreter/packet_response_waiter.rb
+++ b/lib/rex/post/meterpreter/packet_response_waiter.rb
@@ -39,6 +39,9 @@ class PacketResponseWaiter
   # @return [Integer] request ID to wait for
   attr_accessor :rid
 
+  # @return [Boolean] indicates if part of the response has been received
+  attr_accessor :in_progress
+
   #
   # Initializes a response waiter instance for the supplied request
   # identifier.
@@ -46,6 +49,7 @@ class PacketResponseWaiter
   def initialize(rid, completion_routine = nil, completion_param = nil)
     self.rid      = rid.dup
     self.response = nil
+    self.in_progress = false
 
     if (completion_routine)
       self.completion_routine = completion_routine
@@ -69,14 +73,21 @@ class PacketResponseWaiter
   #
   # @param response [Packet]
   # @return [void]
-  def notify(response)
+  def notify(response, in_progress = false)
     if (self.completion_routine)
-      self.response = response
-      self.completion_routine.call(response, self.completion_param)
+      self.in_progress = in_progress
+      unless in_progress
+        self.response = response
+        self.completion_routine.call(response, self.completion_param)
+      end
     else
       self.mutex.synchronize do
-        self.response = response
-        self.cond.signal
+        self.in_progress = in_progress
+        unless in_progress
+          # complete packet, ready for processing...
+          self.response = response
+          self.cond.signal
+        end
       end
     end
   end
@@ -92,7 +103,11 @@ class PacketResponseWaiter
     interval = nil if interval and interval == -1
     self.mutex.synchronize do
       if self.response.nil?
-        self.cond.wait(self.mutex, interval)
+        loop do
+          self.cond.wait(self.mutex, interval)
+          break unless self.in_progress
+          self.in_progress = false
+        end
       end
     end
     return self.response
diff --git a/spec/lib/rex/post/meterpreter/packet_parser_spec.rb b/spec/lib/rex/post/meterpreter/packet_parser_spec.rb
index 1497ebaa9e..22f38ffc5f 100644
--- a/spec/lib/rex/post/meterpreter/packet_parser_spec.rb
+++ b/spec/lib/rex/post/meterpreter/packet_parser_spec.rb
@@ -26,11 +26,12 @@ RSpec.describe Rex::Post::Meterpreter::PacketParser do
 
   it "should parse valid raw data into a packet object" do
     while @raw.length >0
-      parsed_packet = parser.recv(@sock)
+      parsed_packet, in_progress = parser.recv(@sock)
     end
     expect(parsed_packet).to be_a Rex::Post::Meterpreter::Packet
     expect(parsed_packet.type).to eq Rex::Post::Meterpreter::PACKET_TYPE_REQUEST
     expect(parsed_packet.method?("test_method")).to eq true
+    expect(in_progress).to eq false
   end
 
 end