Keep better track of resources

[See rapid7#1623]
[SeeRM rapid7#7692]

Author: egypt

lib/msf/core/exploit/http/server.rb

@@ -36,6 +36,7 @@ def initialize(info = {})
 			], Exploit::Remote::HttpServer
 		)
 
+		@my_resources = []
 		@service_path = nil
 	end
 
@@ -202,6 +203,16 @@ def start_service(opts = {})
 		add_resource(uopts)
 	end
 
+	# Take care of removing any resources that we created
+	def cleanup
+		@my_resources.delete_if do |resource|
+			remove_resource(resource)
+			true
+		end
+
+		super
+	end
+
 	#
 	# Return a Hash containing a best guess at the actual browser and operating
 	# system versions, based on the User-Agent header.
@@ -358,9 +369,15 @@ def report_user_agent(address, request, client_opts={})
 	# NOTE: Calling #add_resource will change the results of subsequent calls
 	# to #get_resource!
 	#
+	# @return (see Rex::Service#add_resource)
 	def add_resource(opts)
 		@service_path = opts['Path']
-		service.add_resource(opts['Path'], opts)
+		res = service.add_resource(opts['Path'], opts)
+
+		# This has to go *after* the call to add_resource in case it raises
+		@my_resources.push(opts['Path'])
+
+		res
 	end
 
 	#

spec/lib/msf/core/exploit/http/server_spec.rb

@@ -0,0 +1,72 @@
+
+# -*- coding:binary -*-
+require 'spec_helper'
+
+require 'msf/core'
+require 'msf/core/exploit/http/server'
+
+describe Msf::Exploit::Remote::HttpServer do
+	subject(:server_module) do
+		mod = Msf::Exploit.allocate
+		mod.extend described_class
+		mod.send(:initialize, {})
+
+		mod
+	end
+
+	let(:mock_service) do
+		mock_service = mock("service")
+		mock_service.stub(:server_name=)
+		mock_service.stub(:add_resource)
+
+		mock_service
+	end
+
+	before do
+		Rex::ServiceManager.stub(:start => mock_service)
+	end
+
+	describe "#add_resource" do
+		it "should call the ServiceManager's add_resource" do
+			server_module.start_service
+
+			mock_service.should_receive(:add_resource)
+			server_module.add_resource('Path' => 'foo')
+		end
+
+		it "should re-raise if the resource has already been added" do
+			server_module.start_service
+
+			mock_service.should_receive(:add_resource).ordered
+			mock_service.should_receive(:add_resource).ordered.and_raise(RuntimeError)
+
+			server_module.add_resource('Path' => 'foo')
+
+			expect { server_module.add_resource('Path' => 'foo') }.to raise_error
+		end
+
+	end
+
+	describe "#cleanup" do
+		it "should not remove resources if none were successfully added" do
+			server_module.should_not_receive(:remove_resource)
+			server_module.cleanup
+		end
+
+		it "should remove successfully-added resources" do
+			# setup
+			server_module.start_service
+			resources = [ 'a', 'b', 'c' ]
+			resources.each { |r| server_module.add_resource('Path' => r) }
+
+			# The service will add one resource as part of #start_service, so
+			# add that to the number that we added manually
+			server_module.should_receive(:remove_resource).exactly(resources.count + 1).times
+			server_module.cleanup
+		end
+
+	end
+
+end
+
+