From 8089ea89a7d70eb0fae09e99d287e8bc54f0a72b Mon Sep 17 00:00:00 2001 From: Andrii Chubatiuk Date: Sat, 5 Oct 2024 01:12:24 +0300 Subject: [PATCH] added custom headers support (#69) * added custom headers support * fixed tests --- lib/fluent/plugin/out_datadog.rb | 12 ++++++++---- test/plugin/test_out_datadog.rb | 12 ++++++------ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/lib/fluent/plugin/out_datadog.rb b/lib/fluent/plugin/out_datadog.rb index 89a3bf8..23c30f0 100644 --- a/lib/fluent/plugin/out_datadog.rb +++ b/lib/fluent/plugin/out_datadog.rb @@ -52,6 +52,7 @@ class RetryableError < StandardError; config_param :max_retries, :integer, :default => -1 config_param :max_backoff, :integer, :default => 30 config_param :use_http, :bool, :default => true + config_param :custom_headers, :hash, :default => {} config_param :use_compression, :bool, :default => true config_param :compression_level, :integer, :default => 6 config_param :no_ssl_validation, :bool, :default => false @@ -98,7 +99,7 @@ def formatted_to_msgpack_binary? def start super - @client = new_client(log, @api_key, @use_http, @use_ssl, @no_ssl_validation, @host, @ssl_port, @port, @http_proxy, @use_compression, @force_v1_routes) + @client = new_client(log, @api_key, @use_http, @use_ssl, @no_ssl_validation, @host, @ssl_port, @port, @http_proxy, @custom_headers, @use_compression, @force_v1_routes) end def shutdown @@ -270,9 +271,9 @@ def gzip_compress(payload, compression_level) end # Build a new transport client - def new_client(logger, api_key, use_http, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, use_compression, force_v1_routes) + def new_client(logger, api_key, use_http, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, custom_headers, use_compression, force_v1_routes) if use_http - DatadogHTTPClient.new logger, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, use_compression, api_key, force_v1_routes + DatadogHTTPClient.new logger, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, custom_headers, use_compression, api_key, force_v1_routes else DatadogTCPClient.new logger, use_ssl, no_ssl_validation, host, ssl_port, port end @@ -310,7 +311,7 @@ class DatadogHTTPClient < DatadogClient require 'net/http' require 'net/http/persistent' - def initialize(logger, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, use_compression, api_key, force_v1_routes = false) + def initialize(logger, use_ssl, no_ssl_validation, host, ssl_port, port, http_proxy, custom_headers, use_compression, api_key, force_v1_routes = false) @logger = logger protocol = use_ssl ? "https" : "http" port = use_ssl ? ssl_port : port @@ -328,6 +329,9 @@ def initialize(logger, use_ssl, no_ssl_validation, host, ssl_port, port, http_pr logger.info("Starting HTTP connection to #{protocol}://#{host}:#{port.to_s} with compression " + (use_compression ? "enabled" : "disabled") + (force_v1_routes ? " using v1 routes" : " using v2 routes")) @client = Net::HTTP::Persistent.new name: "fluent-plugin-datadog-logcollector", proxy: proxy_uri @client.verify_mode = OpenSSL::SSL::VERIFY_NONE if no_ssl_validation + custom_headers.each do |key, value| + @client.override_headers[key] = value + end unless force_v1_routes @client.override_headers["DD-API-KEY"] = api_key @client.override_headers["DD-EVP-ORIGIN"] = "fluent" diff --git a/test/plugin/test_out_datadog.rb b/test/plugin/test_out_datadog.rb index e1cff96..d029abf 100644 --- a/test/plugin/test_out_datadog.rb +++ b/test/plugin/test_out_datadog.rb @@ -229,7 +229,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 500) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key, true + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key, true assert_raise(Fluent::DatadogOutput::RetryableError) do client.send(payload) end @@ -241,7 +241,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 429) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key, true + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key, true assert_raise(Fluent::DatadogOutput::RetryableError) do client.send(payload) end @@ -251,7 +251,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 400) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key, true + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key, true assert_nothing_raised do client.send(payload) end @@ -264,7 +264,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 500, true) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key assert_raise(Fluent::DatadogOutput::RetryableError) do client.send(payload) end @@ -274,7 +274,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 429, true) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key assert_raise(Fluent::DatadogOutput::RetryableError) do client.send(payload) end @@ -284,7 +284,7 @@ def create_valid_subject api_key = 'XXX' stub_dd_request_with_return_code(api_key, 400, true) payload = '{}' - client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, false, api_key + client = Fluent::DatadogOutput::DatadogHTTPClient.new Logger.new(STDOUT), false, false, "datadog.com", 443, 80, nil, {}, false, api_key assert_nothing_raised do client.send(payload) end