Skip to content

Commit 0d9b120

Browse files
committed
Get rid of the suffix
This makes blob cache a little cleaner [FixRM rapid7#7898]
1 parent 21f8e19 commit 0d9b120

File tree

3 files changed

+32
-19
lines changed

3 files changed

+32
-19
lines changed

lib/msf/core/payload.rb

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -503,22 +503,22 @@ def on_session(session)
503503

504504
#
505505
# If the payload has assembly that needs to be compiled, do so now.
506-
# This method takes the raw payload (p), the assembly text (asm), and the
507-
# offsets hash for variables that need to be substituted (off). The suffix
508-
# is used to localize the way the generated payload is cached (whether the
509-
# blob is part of a single, stager, or stage, for example).
510-
#
511-
def build(x, asm, off, suffix = '')
512-
# If there is no assembly to be compiled, then we return a duplicated
513-
# copy of the raw payload blob
506+
#
507+
# Blobs will be cached in the framework's PayloadSet
508+
#
509+
# @see PayloadSet#check_blob_cache
510+
# @param asm [String] Assembly code to be assembled into a raw payload
511+
# @return [String] The final, assembled payload
512+
# @raise ArgumentError if +asm+ is blank
513+
def build(asm, off={})
514514
if(asm.nil? or asm.empty?)
515-
return x.dup
515+
raise ArgumentError, "Assembly must not be empty"
516516
end
517517

518-
# single payloads generate new assembly each time with the options
519-
# substituted in already. It's not appropriate to just use the module
520-
# as the cache key.
521-
cache_key = refname + suffix + asm.hash.to_s
518+
# Use the refname so blobs can be flushed when the module gets
519+
# reloaded and use the hash value to ensure that we're actually
520+
# getting the right blob for the given assembly.
521+
cache_key = refname + asm.hash.to_s
522522
cache_entry = framework.payloads.check_blob_cache(cache_key)
523523

524524
off.each_pair { |option, val|
@@ -576,7 +576,11 @@ def build(x, asm, off, suffix = '')
576576
def internal_generate
577577
# Build the payload, either by using the raw payload blob defined in the
578578
# module or by actually assembling it
579-
raw = build(payload, assembly, offsets, '-stg0')
579+
if assembly and !assembly.empty?
580+
raw = build(assembly, offsets)
581+
else
582+
raw = payload.dup
583+
end
580584

581585
# If the payload is generated and there are offsets to substitute,
582586
# do that now.

lib/msf/core/payload/stager.rb

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,13 +107,19 @@ def encode_stage?
107107
#
108108
# @return [String] The generated payload stage, as a string.
109109
def generate_stage
110+
# XXX: This is nearly identical to Payload#internal_generate
111+
110112
# Compile the stage as necessary
111-
p = build(stage_payload, stage_assembly, stage_offsets, '-stg1')
113+
if stage_assembly and !stage_assembly.empty?
114+
raw = build(stage_assembly, stage_offsets)
115+
else
116+
raw = stage_payload.dup
117+
end
112118

113119
# Substitute variables in the stage
114-
substitute_vars(p, stage_offsets) if (stage_offsets)
120+
substitute_vars(raw, stage_offsets) if (stage_offsets)
115121

116-
return p
122+
return raw
117123
end
118124

119125
#

lib/msf/core/payload_set.rb

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -343,8 +343,11 @@ def stagers
343343
# it must be removed (if one exists)
344344
#
345345
def on_module_reload(mod)
346-
@blob_cache.delete(mod.refname + "-stg0")
347-
@blob_cache.delete(mod.refname + "-stg1")
346+
@blob_cache.each_key do |key|
347+
if key.start_with? mod.refname
348+
@blob_cache.delete(key)
349+
end
350+
end
348351
end
349352

350353
#

0 commit comments

Comments
 (0)