From 1a32930017a3f361ccea43e2fdc8bc396bf28239 Mon Sep 17 00:00:00 2001
From: Brian Cain <bcain@hashicorp.com>
Date: Fri, 7 Dec 2018 10:28:21 -0800
Subject: [PATCH] Add guard_with method for protecting ruby blocks

---
 lib/vagrant/util/experimental.rb            | 26 +++++++++++++++++----
 test/unit/vagrant/util/experimental_test.rb | 23 ++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/lib/vagrant/util/experimental.rb b/lib/vagrant/util/experimental.rb
index 1b9f6d206..cd05cf85a 100644
--- a/lib/vagrant/util/experimental.rb
+++ b/lib/vagrant/util/experimental.rb
@@ -22,20 +22,38 @@ module Vagrant
         # A method for Vagrant internals to determine if a given feature
         # has been abled and can be used.
         #
-        # @param [String] - An array of strings of features to check against
+        # @param [String] feature
         # @return [Boolean] - A hash containing the original array and if it is valid
         def feature_enabled?(feature)
-          experimental = ENV["VAGRANT_EXPERIMENTAL"].to_s.downcase
-          if experimental == "1"
+          experimental = features_requested
+          if experimental.size == 1 && experimental.first == "1"
             return true
           elsif VALID_FEATURES.include?(feature) &&
-                experimental.split(',').include?(feature)
+                experimental.include?(feature)
             return true
           else
             return false
           end
         end
 
+        # Returns the features requested for the experimental flag
+        #
+        # @return [Array] - Returns an array of requested experimental features
+        def features_requested
+          if !defined?(@_requested_features)
+            @_requested_features = ENV["VAGRANT_EXPERIMENTAL"].to_s.downcase.split(',')
+          end
+          @_requested_features
+        end
+
+        # A function to guard experimental blocks of code from being executed
+        #
+        # @param [Array] features - Array of features to guard a method with
+        # @param [Block] block - Block of ruby code to be guarded against
+        def guard_with(*features, &block)
+          yield if block_given? && features.any? {|f| feature_enabled?(f)}
+        end
+
         # @private
         # Reset the cached values for platform. This is not considered a public
         # API and should only be used for testing.
diff --git a/test/unit/vagrant/util/experimental_test.rb b/test/unit/vagrant/util/experimental_test.rb
index 222e19723..7ebe53234 100644
--- a/test/unit/vagrant/util/experimental_test.rb
+++ b/test/unit/vagrant/util/experimental_test.rb
@@ -64,4 +64,27 @@ describe Vagrant::Util::Experimental do
       expect(subject.feature_enabled?("anything")).to eq(false)
     end
   end
+
+  describe "#features_requested" do
+    it "returns an array of requested features" do
+      allow(ENV).to receive(:[]).with("VAGRANT_EXPERIMENTAL").and_return("secret_feature,other_secret")
+      expect(subject.features_requested).to eq(["secret_feature","other_secret"])
+    end
+  end
+
+  describe "#guard_with" do
+    before(:each) do
+      stub_const("Vagrant::Util::Experimental::VALID_FEATURES", ["secret_feature"])
+    end
+
+    it "does not execute the block if the feature is not requested" do
+      allow(ENV).to receive(:[]).with("VAGRANT_EXPERIMENTAL").and_return(nil)
+      expect{|b| subject.guard_with("secret_feature", &b) }.not_to yield_control
+    end
+
+    it "executes the block if the feature is valid and requested" do
+      allow(ENV).to receive(:[]).with("VAGRANT_EXPERIMENTAL").and_return("secret_feature,other_secret")
+      expect{|b| subject.guard_with("secret_feature", &b) }.to yield_control
+    end
+  end
 end