Merge pull request #175 from TheMeier/puppetdb_header

feat: allow to pass additional HTTP headers in puppetdb requests

Author: TheMeier

documentation/bolt_connect_puppetdb.md

@@ -50,6 +50,7 @@ config](configuring_bolt.md) with the following values:
 | --- | --- | --- |
 | `cacert` | `String` | The path to the CA certificate for PuppetDB. |
 | `connect_timeout` | `Integer` | How long to wait in seconds when establishing connections with PuppetDB. |
+| `headers` | `Hash` | A map of HTTP headers to add to PuppetDB requests. |
 | `read_timeout` | `Integer` | How long to wait in seconds for a response from PuppetDB. |
 | `server_urls` | `Array` | An array of strings containing the PuppetDB host to connect to. Include the protocol `https` and the port, which is usually `8081`. For example, `https://my-puppetdb-server.example.com:8081`. The Bolt PuppetDB client attempts to connect to each host in the list until it makes a successful connection. |
 
@@ -94,6 +95,16 @@ puppetdb:
   token: ~/.puppetlabs/token
 ```
 
+To use custom headers, such as for OAuth authentication:
+
+```
+puppetdb:
+  server_urls: ["https://puppet.example.com:8081"]
+  cacert: /etc/puppetlabs/puppet/ssl/certs/ca.pem
+  headers:
+    Authorization: "Bearer <token-content>"
+```
+
 ## Configuring multiple PuppetDB instances
 
 The Bolt PuppetDB Client supports connections to multiple PuppetDB instances. To
@@ -244,15 +255,15 @@ plan puppetdb_query_targets {
   # this returns an array of objects, each object containing a "certname" parameter:
   # [ {"certname": "node1"}, {"certname": "node2"} ]
   $query_results = puppetdb_query("nodes[certname] {}")
-  
+
   # since puppetdb_query() returns the JSON results from the API call, we need to transform this
   # data into Targets to use it in one of the run_*() functions.
   # extract the "certname" values, so now we have an array of hostnames
   $certnames = $query_results.map |$r| { $r['certname'] }
-  
+
   # transform the arary of certnames into an array of Targets
   $targets = get_targets($certnames)
-  
+
   # gather facts about all of the nodes
   run_task('facts', $targets)
 }
@@ -270,10 +281,10 @@ plan puppetdb_plugin_targets {
     'query'   => 'nodes[certname] {}',
   }
   $references = resolve_references($refs)
-  
+
   # maps the results into a list of Target objects
   $targets = $references.map |$r| { Target.new($r) }
-  
+
   # gather facts about all of the nodes
   run_task('facts', $targets)
 }

lib/bolt/config/options.rb

@@ -76,6 +76,12 @@ module Options
           _example: 120,
           _plugin: true
         },
+        "headers" => {
+          description: "A map of HTTP headers to add to PuppetDB requests.",
+          type: Hash,
+          _example: { "Authorization" => "Bearer <token>" },
+          _plugin: true
+        },
         "key" => {
           description: "The private key for the certificate.",
           type: String,

lib/bolt/puppetdb/config.rb

@@ -149,6 +149,14 @@ def key
         @settings['key']
       end
 
+      def headers
+        if @settings['headers'] && !@settings['headers'].is_a?(Hash)
+          raise Bolt::PuppetDBError, "headers must be a Hash"
+        end
+
+        @settings['headers']
+      end
+
       def validate_cert_and_key
         if (@settings['cert'] && !@settings['key']) ||
            (!@settings['cert'] && @settings['key'])

lib/bolt/puppetdb/instance.rb

@@ -139,6 +139,7 @@ def uri
 
       def headers
         headers = { 'Content-Type' => 'application/json' }
+        headers.merge!(@config.headers) if @config.headers
         headers['X-Authentication'] = @config.token if @config.token
         headers
       end

schemas/bolt-defaults.schema.json

@@ -299,6 +299,17 @@
                 }
               ]
             },
+            "headers": {
+              "description": "A map of HTTP headers to add to PuppetDB requests.",
+              "oneOf": [
+                {
+                  "type": "object"
+                },
+                {
+                  "$ref": "#/definitions/_plugin"
+                }
+              ]
+            },
             "key": {
               "description": "The private key for the certificate.",
               "oneOf": [
@@ -395,6 +406,17 @@
                       }
                     ]
                   },
+                  "headers": {
+                    "description": "A map of HTTP headers to add to PuppetDB requests.",
+                    "oneOf": [
+                      {
+                        "type": "object"
+                      },
+                      {
+                        "$ref": "#/definitions/_plugin"
+                      }
+                    ]
+                  },
                   "key": {
                     "description": "The private key for the certificate.",
                     "oneOf": [

schemas/bolt-project.schema.json

@@ -429,6 +429,17 @@
                 }
               ]
             },
+            "headers": {
+              "description": "A map of HTTP headers to add to PuppetDB requests.",
+              "oneOf": [
+                {
+                  "type": "object"
+                },
+                {
+                  "$ref": "#/definitions/_plugin"
+                }
+              ]
+            },
             "key": {
               "description": "The private key for the certificate.",
               "oneOf": [
@@ -525,6 +536,17 @@
                       }
                     ]
                   },
+                  "headers": {
+                    "description": "A map of HTTP headers to add to PuppetDB requests.",
+                    "oneOf": [
+                      {
+                        "type": "object"
+                      },
+                      {
+                        "$ref": "#/definitions/_plugin"
+                      }
+                    ]
+                  },
                   "key": {
                     "description": "The private key for the certificate.",
                     "oneOf": [

spec/unit/puppetdb/config_spec.rb

@@ -254,4 +254,25 @@
       Bolt::PuppetDB::Config.new(config: {}, load_defaults: true)
     end
   end
+
+  context 'when validating headers' do
+    let(:options) { { 'server_urls' => ['https://puppetdb:8081'], 'headers' => headers } }
+    let(:config) { Bolt::PuppetDB::Config.new(config: options) }
+
+    context 'with valid headers' do
+      let(:headers) { { 'Authorization' => 'Bearer token' } }
+
+      it 'returns the headers' do
+        expect(config.headers).to eq(headers)
+      end
+    end
+
+    context 'with invalid headers' do
+      let(:headers) { 'Authorization: Bearer token' }
+
+      it 'raises an error' do
+        expect { config.headers }.to raise_error(Bolt::PuppetDBError, "headers must be a Hash")
+      end
+    end
+  end
 end

spec/unit/puppetdb/instance_spec.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'bolt/puppetdb/instance'
+
+describe Bolt::PuppetDB::Instance do
+  let(:token) { 'token' }
+  let(:config) do
+    {
+      'server_urls' => ["https://puppet.example.com:8081"],
+      'cacert'      => '/etc/puppetlabs/puppet/ssl/certs/ca.pem',
+      'token'       => '~/.puppetlabs/token'
+    }
+  end
+  let(:instance) { described_class.new(config: config) }
+
+  before(:each) do
+    allow(File).to receive(:exist?).and_return(true)
+    allow(File).to receive(:read).and_call_original
+    allow(File).to receive(:read).with(File.expand_path('~/.puppetlabs/token')).and_return(token)
+  end
+
+  context "#headers" do
+    it "includes Content-Type" do
+      expect(instance.headers).to include('Content-Type' => 'application/json')
+    end
+
+    it "includes X-Authentication token" do
+      expect(instance.headers).to include('X-Authentication' => token)
+    end
+
+    context "with custom headers" do
+      let(:config) do
+        {
+          'server_urls' => ["https://puppet.example.com:8081"],
+          'headers'     => { 'Authorization' => 'Bearer info' }
+        }
+      end
+
+      it "includes custom headers" do
+        expect(instance.headers).to include('Authorization' => 'Bearer info')
+      end
+
+      it "does not include X-Authentication if no token" do
+        # config does not have 'token', so it falls back to DEFAULT_TOKEN.
+        # We need to simulate no default token file.
+        allow(File).to receive(:exist?).with(Bolt::PuppetDB::Config::DEFAULT_TOKEN).and_return(false)
+        expect(instance.headers).not_to have_key('X-Authentication')
+      end
+    end
+
+    context "with custom headers overlapping Content-Type" do
+      let(:config) do
+        {
+          'server_urls' => ["https://puppet.example.com:8081"],
+          'headers'     => { 'Content-Type' => 'application/x-yaml' }
+        }
+      end
+
+      it "overrides default Content-Type" do
+        expect(instance.headers).to include('Content-Type' => 'application/x-yaml')
+      end
+    end
+  end
+end