Merge pull request #1520 from higanworks/imds-v2

Signed-off-by: Tim Smith <tsmith@chef.io>
chef · Jan 25, 2021 · c33e81b · c33e81b
2 parents f4e5b68 + dc6955e
commit c33e81b
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 69 deletions.
diff --git a/lib/ohai/mixin/ec2_metadata.rb b/lib/ohai/mixin/ec2_metadata.rb
@@ -51,7 +51,7 @@ module Ec2Metadata
       def best_api_version
         @api_version ||= begin
           logger.trace("Mixin EC2: Fetching http://#{EC2_METADATA_ADDR}/ to determine the latest supported metadata release")
-          response = http_client.get("/")
+          response = http_client.get("/", { 'X-aws-ec2-metadata-token': v2_token })
           if response.code == "404"
             logger.trace("Mixin EC2: Received HTTP 404 from metadata server while determining API version, assuming 'latest'")
             return "latest"
@@ -84,6 +84,10 @@ def http_client
         end
       end
 
+      def v2_token
+        @v2_token ||= http_client.put("/latest/api/token", nil, { 'X-aws-ec2-metadata-token-ttl-seconds': "60" })&.body
+      end
+
       # Get metadata for a given path and API version
       #
       # Typically, a 200 response is expected for valid metadata.
@@ -93,7 +97,7 @@ def http_client
       def metadata_get(id, api_version)
         path = "/#{api_version}/meta-data/#{id}"
         logger.trace("Mixin EC2: Fetching http://#{EC2_METADATA_ADDR}#{path}")
-        response = http_client.get(path)
+        response = http_client.get(path, { 'X-aws-ec2-metadata-token': v2_token })
         case response.code
         when "200"
           response.body
@@ -174,13 +178,13 @@ def fetch_json_dir_metadata(id, api_version)
 
       def fetch_userdata
         logger.trace("Mixin EC2: Fetching http://#{EC2_METADATA_ADDR}/#{best_api_version}/user-data/")
-        response = http_client.get("/#{best_api_version}/user-data/")
+        response = http_client.get("/#{best_api_version}/user-data/", { 'X-aws-ec2-metadata-token': v2_token })
         response.code == "200" ? response.body : nil
       end
 
       def fetch_dynamic_data
         @fetch_dynamic_data ||= begin
-          response = http_client.get("/#{best_api_version}/dynamic/instance-identity/document/")
+          response = http_client.get("/#{best_api_version}/dynamic/instance-identity/document/", { 'X-aws-ec2-metadata-token': v2_token })
 
           if json?(response.body) && response.code == "200"
             FFI_Yajl::Parser.parse(response.body)

diff --git a/spec/unit/mixin/ec2_metadata_spec.rb b/spec/unit/mixin/ec2_metadata_spec.rb
@@ -15,14 +15,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
 require "spec_helper"
 require "ohai/mixin/ec2_metadata"
 
 describe Ohai::Mixin::Ec2Metadata do
   let(:mixin) do
     metadata_object = Object.new.extend(described_class)
     http_client = double("Net::HTTP client")
+    allow(http_client).to receive(:put) { double("Net::HTTP::PUT Response", body: "AQAEAE4UUd-3NE5EEeYYXKxicVfDOHsx0YSHFFSuCvo2GfCcxzJsvg==", code: "200") }
     allow(http_client).to receive(:get).and_return(response)
     allow(metadata_object).to receive(:http_client).and_return(http_client)
     metadata_object
@@ -36,7 +36,6 @@
   describe "#best_api_version" do
     context "with a sorted list of metadata versions" do
       let(:response) { double("Net::HTTP Response", body: "1.0\n2011-05-01\n2012-01-12\nUnsupported", code: "200") }
-
       it "returns the most recent version" do
         expect(mixin.best_api_version).to eq("2012-01-12")
       end
@@ -74,6 +73,13 @@
         expect { mixin.best_api_version }.to raise_error(RuntimeError)
       end
     end
+
+    context "when metadata service is disabled" do
+      let(:response) { double("Net::HTTP::PUT Response", body: "403 - Forbidden", code: "403") }
+      it "raises an error" do
+        expect { mixin.best_api_version }.to raise_error(RuntimeError)
+      end
+    end
   end
 
   describe "#metadata_get" do