Pull out port collision detection/correction into the CheckPortCollision middleware

lib/vagrant/action/builder.rb

@@ -38,10 +38,12 @@ module Vagrant
       #
       # @param [Class] middleware The middleware class
       def use(middleware, *args, &block)
-        if middleware.kind_of?(Builder)
-          # Prepend with a environment setter if args are given
-          self.use(Env::Set, *args, &block) if !args.empty? && args.first.is_a?(Hash)
+        # Prepend with a environment setter if args are given
+        if !args.empty? && args.first.is_a?(Hash) && middleware != Env::Set
+          self.use(Env::Set, args.shift, &block)
+        end
 
+        if middleware.kind_of?(Builder)
           # Merge in the other builder's stack into our own
           self.stack.concat(middleware.stack)
         else

lib/vagrant/action/builtin.rb

@@ -20,6 +20,7 @@ module Vagrant
         Builder.new do
           use General::Validate
           use VM::CheckAccessible
+          use VM::CheckPortCollisions, :port_collision_handler => :correct
           use VM::CleanMachineFolder
           use VM::ClearForwardedPorts
           use VM::ForwardPorts

lib/vagrant/action/env/set.rb

@@ -4,12 +4,15 @@ module Vagrant
       # A middleware which just sets up the environment with some
       # options which are passed to it.
       class Set
-        def initialize(app,env,options=nil)
-          @app = app
-          env.merge!(options || {})
+        def initialize(app, env, options=nil)
+          @app     = app
+          @options = options || {}
         end
 
         def call(env)
+          # Merge the options that were given to us
+          env.merge!(@options)
+
           @app.call(env)
         end
       end

lib/vagrant/action/vm/check_port_collisions.rb

@@ -9,17 +9,62 @@ module Vagrant
         end
 
         def call(env)
-          existing = env[:vm].driver.read_used_ports
+          # For the handlers...
+          @env = env
 
+          # Figure out how we handle port collisions. By default we error.
+          handler  = env[:port_collision_handler] || :error
+
+          existing = env[:vm].driver.read_used_ports
           env[:vm].config.vm.forwarded_ports.each do |name, options|
             if existing.include?(options[:hostport].to_i)
-              # We have a collision!
-              raise Errors::ForwardPortCollisionResume
+              # We have a collision! Handle it
+              send("handle_#{handler}".to_sym, name, options, existing)
             end
           end
 
           @app.call(env)
         end
+
+        # Handles a port collision by raising an exception.
+        def handle_error(name, options, existing_ports)
+          raise Errors::ForwardPortCollisionResume
+        end
+
+        # Handles a port collision by attempting to fix it.
+        def handle_correct(name, options, existing_ports)
+          if !options[:auto]
+            # Auto fixing is disabled for this port forward, so we
+            # must throw an error so the user can fix it.
+            raise Errors::ForwardPortCollision, :name => name,
+                                                :host_port => options[:hostport].to_s,
+                                                :guest_port => options[:guestport].to_s
+          end
+
+          # Get the auto port range and get rid of the used ports and
+          # ports which are being used in other forwards so we're just
+          # left with available ports.
+          range = @env[:vm].config.vm.auto_port_range.to_a
+          range -= @env[:vm].config.vm.forwarded_ports.collect { |n, o| o[:hostport].to_i }
+          range -= existing_ports
+
+          if range.empty?
+            raise Errors::ForwardPortAutolistEmpty, :vm_name => @env[:vm].name,
+                                                    :name => name,
+                                                    :host_port => options[:hostport].to_s,
+                                                    :guest_port => options[:guestport].to_s
+          end
+
+          # Set the port up to be the first one and add that port to
+          # the used list.
+          options[:hostport] = range.shift
+          existing_ports << options[:hostport]
+
+          # Notify the user
+          @env[:ui].info(I18n.t("vagrant.actions.vm.forward_ports.fixed_collision",
+                                :name => name,
+                                :new_port => options[:hostport]))
+        end
       end
     end
   end

lib/vagrant/action/vm/forward_ports.rb

@@ -7,7 +7,6 @@ module Vagrant
           @env = env
 
           threshold_check
-          external_collision_check
         end
 
         #--------------------------------------------------------------
@@ -26,56 +25,6 @@ module Vagrant
           end
         end
 
-        # This method checks for any port collisions with any VMs
-        # which are already created (by Vagrant or otherwise).
-        # report the collisions detected or will attempt to fix them
-        # automatically if the port is configured to do so.
-        def external_collision_check
-          existing = @env[:vm].driver.read_used_ports
-          @env[:vm].config.vm.forwarded_ports.each do |name, options|
-            if existing.include?(options[:hostport].to_i)
-              handle_collision(name, options, existing)
-            end
-          end
-        end
-
-        # Handles any collisions. This method will either attempt to
-        # fix the collision automatically or will raise an error if
-        # auto fixing is disabled.
-        def handle_collision(name, options, existing_ports)
-          if !options[:auto]
-            # Auto fixing is disabled for this port forward, so we
-            # must throw an error so the user can fix it.
-            raise Errors::ForwardPortCollision, :name => name,
-                                                :host_port => options[:hostport].to_s,
-                                                :guest_port => options[:guestport].to_s
-          end
-
-          # Get the auto port range and get rid of the used ports and
-          # ports which are being used in other forwards so we're just
-          # left with available ports.
-          range = @env[:vm].config.vm.auto_port_range.to_a
-          range -= @env[:vm].config.vm.forwarded_ports.collect { |n, o| o[:hostport].to_i }
-          range -= existing_ports
-
-          if range.empty?
-            raise Errors::ForwardPortAutolistEmpty, :vm_name => @env[:vm].name,
-                                                    :name => name,
-                                                    :host_port => options[:hostport].to_s,
-                                                    :guest_port => options[:guestport].to_s
-          end
-
-          # Set the port up to be the first one and add that port to
-          # the used list.
-          options[:hostport] = range.shift
-          existing_ports << options[:hostport]
-
-          # Notify the user
-          @env[:ui].info(I18n.t("vagrant.actions.vm.forward_ports.fixed_collision",
-                                :name => name,
-                                :new_port => options[:hostport]))
-        end
-
         #--------------------------------------------------------------
         # Execution
         #--------------------------------------------------------------

test/unit/vagrant/action/builder_test.rb

@@ -51,21 +51,16 @@ describe Vagrant::Action::Builder do
       data[:one].should == true
     end
 
-    it "should be able to set additional variables if using another builder" do
+    it "should be able to set additional variables when using" do
       data  = { }
       proc1 = Proc.new { |env| env[:data] += 1 }
 
       # Build the first builder
       one   = described_class.new
-      one.use proc1
+      one.use proc1, :data => 5
+      one.call(data)
 
-      # Add it to this builder
-      two   = described_class.new
-      two.use one, :data => 10
-
-      # Call the 2nd and verify results
-      two.call(data)
-      data[:data].should == 11
+      data[:data].should == 6
     end
   end