feat: allow to pass additional HTTP headers in puppetdb requests

This can be useful if you want to access PuppetDB via custom authentication.
Eg if you want to access PuppetDB behind an API gateway that requires oauth tokens
or other custom headers.

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