From 52b0d489d7a5523605d689e49b4c990d992e627b Mon Sep 17 00:00:00 2001
From: Matt Brictson <matt@mattbrictson.com>
Date: Wed, 1 Jul 2015 17:24:26 -0700
Subject: [PATCH] Prevent enhance of load:defaults after invoke

Print a warning and abort if "load:defaults" is erroneously invoked after
capistrano is already loaded, e.g. when a plugin is loaded in `deploy.rb`
instead of `Capfile`.
---
 CHANGELOG.md                               |  3 +++
 lib/capistrano/immutable_task.rb           | 29 ++++++++++++++++++++++
 lib/capistrano/setup.rb                    |  2 ++
 spec/lib/capistrano/immutable_task_spec.rb | 26 +++++++++++++++++++
 4 files changed, 60 insertions(+)
 create mode 100644 lib/capistrano/immutable_task.rb
 create mode 100644 spec/lib/capistrano/immutable_task_spec.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 22e36cfc..b0c9e118 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,9 @@ https://github.com/capistrano/capistrano/compare/v3.4.0...HEAD
   * Allow after() to refer to tasks that have not been loaded yet (@jcoglan)
   * Ensure scm fetch_revision methods strip trailing whitespace (@mattbrictson)
   * Allow use "all" as string for server filtering (@theist)
+  * Print a warning and abort if "load:defaults" is erroneously invoked after
+    capistrano is already loaded, e.g. when a plugin is loaded in `deploy.rb`
+    instead of `Capfile`. (@mattbrictson)
 
 ## `3.4.0`
 
diff --git a/lib/capistrano/immutable_task.rb b/lib/capistrano/immutable_task.rb
new file mode 100644
index 00000000..0e86f834
--- /dev/null
+++ b/lib/capistrano/immutable_task.rb
@@ -0,0 +1,29 @@
+module Capistrano
+  # This module extends a Rake::Task to freeze it to prevent it from being
+  # enhanced. This is used to prevent users from enhancing a task at the wrong
+  # point of Capistrano's boot process, which can happen if a Capistrano plugin
+  # is loaded in deploy.rb by mistake (instead of in the Capfile).
+  #
+  # Usage:
+  #
+  # task = Rake.application["load:defaults"]
+  # task.invoke
+  # task.extend(Capistrano::ImmutableTask) # prevent further modifications
+  #
+  module ImmutableTask
+    def self.extended(task)
+      task.freeze
+    end
+
+    def enhance(*args, &block)
+      $stderr.puts <<-MESSAGE
+WARNING: #{name} has already been invoked and can no longer be modified.
+Check that you haven't loaded a Capistrano plugin in deploy.rb by mistake.
+Plugins must be loaded in the Capfile to initialize properly.
+MESSAGE
+
+      # This will raise a frozen object error
+      super(*args, &block)
+    end
+  end
+end
diff --git a/lib/capistrano/setup.rb b/lib/capistrano/setup.rb
index c13bddad..01bd5fe3 100644
--- a/lib/capistrano/setup.rb
+++ b/lib/capistrano/setup.rb
@@ -1,3 +1,4 @@
+require "capistrano/immutable_task"
 include Capistrano::DSL
 
 namespace :load do
@@ -11,6 +12,7 @@ stages.each do |stage|
     set(:stage, stage.to_sym)
 
     invoke 'load:defaults'
+    Rake.application["load:defaults"].extend(Capistrano::ImmutableTask)
     load deploy_config_path
     load stage_config_path.join("#{stage}.rb")
     load "capistrano/#{fetch(:scm)}.rb"
diff --git a/spec/lib/capistrano/immutable_task_spec.rb b/spec/lib/capistrano/immutable_task_spec.rb
new file mode 100644
index 00000000..1f246dc0
--- /dev/null
+++ b/spec/lib/capistrano/immutable_task_spec.rb
@@ -0,0 +1,26 @@
+require 'spec_helper'
+require 'rake'
+require 'capistrano/immutable_task'
+
+module Capistrano
+  describe ImmutableTask do
+    it 'prints warning and raises when task is enhanced' do
+      extend(Rake::DSL)
+
+      load_defaults = Rake::Task.define_task('load:defaults')
+      load_defaults.extend(Capistrano::ImmutableTask)
+
+      $stderr.expects(:puts).with do |message|
+        message =~ /^WARNING: load:defaults has already been invoked/
+      end
+
+      expect {
+        namespace :load do
+          task :defaults do
+            # Never reached since load_defaults is frozen and can't be enhanced
+          end
+        end
+      }.to raise_error(/frozen/i)
+    end
+  end
+end