From 58f2de887a74f7230527b9c5a9fdd937e104cf2d Mon Sep 17 00:00:00 2001
From: donoghuc <cas.donoghue@gmail.com>
Date: Mon, 8 Apr 2024 14:38:21 -0700
Subject: [PATCH] (GH-3296) Try both token and cert based auth for puppetdb

Previously token was always included in header. With this commit, try POST requests with token in header. If that returns 401 AND cert/key is configured, retry with auth header excluded.

!feature

* **Try cert based auth, failback to cert based auth for puppetdb** ([#3296](#3296))

  When both a token and cert are specified, try using the token in the
  x-auth header for puppetdb POSTs. If that responds with a 401, try
  again with the token excluded.
---
 lib/bolt/puppetdb/instance.rb | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/lib/bolt/puppetdb/instance.rb b/lib/bolt/puppetdb/instance.rb
index 15e927c8ce..baf6516687 100644
--- a/lib/bolt/puppetdb/instance.rb
+++ b/lib/bolt/puppetdb/instance.rb
@@ -16,6 +16,19 @@ def initialize(config:, project: nil, load_defaults: false)
         @logger      = Bolt::Logger.logger(self)
       end
 
+      def post_puppetdb(url, body)
+        response = http_client.post(url, body: body, header: headers(@config.token))
+        if response.status == 401 && @token_and_cert
+          @logger.debug("Invalid token: #{response.body}, retrying with cert based auth")
+          response = http_client.post(url, body: body, header: headers)
+          if response.ok?
+            @logger.debug("Puppetdb token is invalid, but certs are not. No longer including token.")
+            @bad_token = true
+          end
+        end
+        response
+      end
+
       def make_query(query, path = nil)
         body = JSON.generate(query: query)
         url = "#{uri}/pdb/query/v4"
@@ -23,7 +36,7 @@ def make_query(query, path = nil)
 
         begin
           @logger.debug("Sending PuppetDB query to #{url}")
-          response = http_client.post(url, body: body, header: headers)
+          response = post_puppetdb(url, body)
         rescue StandardError => e
           raise Bolt::PuppetDBFailoverError, "Failed to query PuppetDB: #{e}"
         end
@@ -81,7 +94,7 @@ def send_command(command, version, payload)
         # Send the command to PDB
         begin
           @logger.debug("Sending PuppetDB command '#{command}' to #{url}")
-          response = http_client.post(url.to_s, body: body, header: headers)
+          response = post_puppetdb(url.to_s, body)
         rescue StandardError => e
           raise Bolt::PuppetDBFailoverError, "Failed to invoke PuppetDB command: #{e}"
         end
@@ -109,11 +122,15 @@ def http_client
         require 'httpclient'
         @logger.trace("Creating HTTP Client")
         @http = HTTPClient.new
-        @http.ssl_config.set_client_cert_file(@config.cert, @config.key) if @config.cert
         @http.ssl_config.add_trust_ca(@config.cacert)
         @http.connect_timeout = @config.connect_timeout if @config.connect_timeout
         @http.receive_timeout = @config.read_timeout if @config.read_timeout
-
+        # Determine if there are both token and cert auth methods defined
+        @token_and_cert = false
+        if @config.cert
+          @http.ssl_config.set_client_cert_file(@config.cert, @config.key)
+          @token_and_cert = !@config.token.nil?
+        end
         @http
       end
 
@@ -136,9 +153,9 @@ def uri
         uri
       end
 
-      def headers
+      def headers(token = nil)
         headers = { 'Content-Type' => 'application/json' }
-        headers['X-Authentication'] = @config.token if @config.token
+        headers['X-Authentication'] = token if token && !@bad_token
         headers
       end
     end