Land rapid7#1776, assembly payload blob cache fix

Author: egypt

lib/msf/core/payload.rb

@@ -503,19 +503,22 @@ def on_session(session)
 
 	#
 	# If the payload has assembly that needs to be compiled, do so now.
-	# This method takes the raw payload (p), the assembly text (asm), and the
-	# offsets hash for variables that need to be substituted (off).  The suffix
-	# is used to localize the way the generated payload is cached (whether the
-	# blob is part of a single, stager, or stage, for example).
-	#
-	def build(x, asm, off, suffix = '')
-		# If there is no assembly to be compiled, then we return a duplicated
-		# copy of the raw payload blob
+	#
+	# Blobs will be cached in the framework's PayloadSet
+	#
+	# @see PayloadSet#check_blob_cache
+	# @param asm [String] Assembly code to be assembled into a raw payload
+	# @return [String] The final, assembled payload
+	# @raise ArgumentError if +asm+ is blank
+	def build(asm, off={})
 		if(asm.nil? or asm.empty?)
-			return x.dup
+			raise ArgumentError, "Assembly must not be empty"
 		end
 
-		cache_key   = refname + suffix
+		# Use the refname so blobs can be flushed when the module gets
+		# reloaded and use the hash value to ensure that we're actually
+		# getting the right blob for the given assembly.
+		cache_key   = refname + asm.hash.to_s
 		cache_entry = framework.payloads.check_blob_cache(cache_key)
 
 		off.each_pair { |option, val|
@@ -573,7 +576,11 @@ def build(x, asm, off, suffix = '')
 	def internal_generate
 		# Build the payload, either by using the raw payload blob defined in the
 		# module or by actually assembling it
-		raw = build(payload, assembly, offsets, '-stg0')
+		if assembly and !assembly.empty?
+			raw = build(assembly, offsets)
+		else
+			raw = payload.dup
+		end
 
 		# If the payload is generated and there are offsets to substitute,
 		# do that now.

lib/msf/core/payload/stager.rb

@@ -107,13 +107,19 @@ def encode_stage?
 	#
 	# @return [String] The generated payload stage, as a string.
 	def generate_stage
+		# XXX: This is nearly identical to Payload#internal_generate
+
 		# Compile the stage as necessary
-		p = build(stage_payload, stage_assembly, stage_offsets, '-stg1')
+		if stage_assembly and !stage_assembly.empty?
+			raw = build(stage_assembly, stage_offsets)
+		else
+			raw = stage_payload.dup
+		end
 
 		# Substitute variables in the stage
-		substitute_vars(p, stage_offsets) if (stage_offsets)
+		substitute_vars(raw, stage_offsets) if (stage_offsets)
 
-		return p
+		return raw
 	end
 
 	#

lib/msf/core/payload_set.rb

@@ -343,8 +343,11 @@ def stagers
 	# it must be removed (if one exists)
 	#
 	def on_module_reload(mod)
-		@blob_cache.delete(mod.refname + "-stg0")
-		@blob_cache.delete(mod.refname + "-stg1")
+		@blob_cache.each_key do |key|
+			if key.start_with? mod.refname
+				@blob_cache.delete(key)
+			end
+		end
 	end
 
 	#