fluent-plugin-azure-logs-ingestion 0.1.0

data/lib/fluent/plugin/azure_logs_ingestion/client.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'net/http'
+require 'uri'
+require 'openssl'
+require 'securerandom'
+require 'time'
+
+module Fluent
+  module Plugin
+    module AzureLogsIngestion
+      class Client
+        def initialize(endpoint:, dcr_immutable_id:, stream_name:, logger:)
+          @endpoint = endpoint
+          @dcr_immutable_id = dcr_immutable_id
+          @stream_name = stream_name
+          @log = logger
+        end
+
+        def send_payload(payload:, bearer_token:)
+          uri = build_uri
+          @log.debug('sending logs ingestion request', uri: uri.to_s, content_length: payload.content_length, content_encoding: payload.content_encoding)
+          request = Net::HTTP::Post.new(uri)
+          request['Authorization'] = "Bearer #{bearer_token}"
+          request['Content-Type'] = 'application/json'
+          request['Content-Encoding'] = payload.content_encoding if payload.content_encoding
+          request['x-ms-client-request-id'] = SecureRandom.uuid
+          request.body_stream = payload.io
+          request.content_length = payload.content_length
+
+          response = perform_request(uri, request)
+          handle_response(response)
+          true
+        rescue Timeout::Error, Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::ECONNRESET, EOFError, SocketError, IOError, SystemCallError => error
+          raise "logs ingestion request failed: #{error.message}"
+        ensure
+          payload.io.rewind if payload&.io
+        end
+
+        private
+
+        def build_uri
+          base = @endpoint.end_with?('/') ? @endpoint : "#{@endpoint}/"
+          URI.join(base, "dataCollectionRules/#{@dcr_immutable_id}/streams/#{@stream_name}?api-version=2023-01-01")
+        end
+
+        def perform_request(uri, request)
+          http = Net::HTTP.new(uri.host, uri.port)
+          http.use_ssl = uri.scheme == 'https'
+          http.open_timeout = 10
+          http.read_timeout = 60
+          http.verify_mode = OpenSSL::SSL::VERIFY_PEER if http.use_ssl?
+          http.request(request)
+        end
+
+        def handle_response(response)
+          code = response.code.to_i
+          @log.debug('received logs ingestion response', code: code, message: response.message, retry_after: response['Retry-After'])
+          return handle_success_response(response) if response.is_a?(Net::HTTPSuccess)
+
+          message = "#{response.code} #{response.message} #{String(response.body).strip}".strip
+          case code
+          when 400, 401, 403, 413
+            raise Fluent::UnrecoverableError, message
+          when 429
+            @log.warn('received retryable 429 from Logs Ingestion API', retry_after: response['Retry-After']) if response['Retry-After']
+            raise message
+          when 500..599
+            raise message
+          else
+            if code >= 400 && code < 500
+              raise Fluent::UnrecoverableError, message
+            end
+
+            raise message
+          end
+        end
+
+        def handle_success_response(response)
+          body = String(response.body)
+          return true if body.empty?
+
+          JSON.parse(body)
+          true
+        rescue JSON::ParserError => error
+          raise "successful response body was not valid JSON: #{error.message}"
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/azure_logs_ingestion/payload_builder.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'tempfile'
+require 'zlib'
+
+module Fluent
+  module Plugin
+    module AzureLogsIngestion
+      class PayloadBuilder
+        MAX_BYTES = 1_048_576
+        Result = Struct.new(
+          :io,
+          :content_encoding,
+          :content_length,
+          :raw_size,
+          :gzip_size,
+          :record_count,
+          keyword_init: true
+        ) do
+          def close!
+            io.close!
+          rescue StandardError
+            nil
+          end
+        end
+
+        def initialize(gzip:)
+          @gzip = gzip
+        end
+
+        def build(chunk)
+          raw_file = Tempfile.new('azure-logs-ingestion-raw')
+          raw_file.binmode
+          gzip_file = nil
+
+          raw_size = 0
+          record_count = 0
+          first_record = true
+
+          raw_size += write_bytes(raw_file, '[')
+          chunk.each do |_event_time, record|
+            json = JSON.generate(record.dup)
+            raw_size += write_bytes(raw_file, ',') unless first_record
+            raw_size += write_bytes(raw_file, json)
+            first_record = false
+            record_count += 1
+          end
+          raw_size += write_bytes(raw_file, ']')
+          raw_file.flush
+          raw_file.rewind
+
+          gzip_size = nil
+          io = raw_file
+          content_encoding = nil
+          content_length = raw_size
+
+          if @gzip
+            gzip_file, gzip_size = gzip_file_from(raw_file)
+            io = gzip_file
+            content_encoding = 'gzip'
+            content_length = gzip_size
+            raw_file.close!
+          end
+
+          validate!(raw_size: raw_size, gzip_size: gzip_size)
+
+          Result.new(
+            io: io,
+            content_encoding: content_encoding,
+            content_length: content_length,
+            raw_size: raw_size,
+            gzip_size: gzip_size,
+            record_count: record_count
+          )
+        rescue StandardError
+          gzip_file.close! if gzip_file
+          raw_file.close! if raw_file
+          raise
+        end
+
+        private
+
+        def validate!(raw_size:, gzip_size:)
+          raise Fluent::UnrecoverableError, "payload size #{raw_size} exceeds #{MAX_BYTES} bytes" if raw_size > MAX_BYTES
+          if gzip_size && gzip_size > MAX_BYTES
+            raise Fluent::UnrecoverableError, "gzip payload size #{gzip_size} exceeds #{MAX_BYTES} bytes"
+          end
+        end
+
+        def gzip_file_from(source_file)
+          gzip_file = Tempfile.new('azure-logs-ingestion-gzip')
+          gzip_file.binmode
+          Zlib::GzipWriter.open(gzip_file.path) do |writer|
+            source_file.rewind
+            IO.copy_stream(source_file, writer)
+          end
+          gzip_file.close
+          gzip_file.open
+          gzip_file.binmode
+          gzip_file.rewind
+          [gzip_file, File.size(gzip_file.path)]
+        end
+
+        def write_bytes(io, string)
+          io.write(string)
+          string.bytesize
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/azure_logs_ingestion/version.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Fluent
+  module Plugin
+    module AzureLogsIngestion
+      VERSION = '0.1.0'
+    end
+  end
+end

data/lib/fluent/plugin/out_azure_logs_ingestion.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require 'fluent/plugin/output'
+require_relative 'azure_logs_ingestion/version'
+require_relative 'azure_logs_ingestion/auth'
+require_relative 'azure_logs_ingestion/payload_builder'
+require_relative 'azure_logs_ingestion/client'
+
+module Fluent
+  module Plugin
+    class AzureLogsIngestionOutput < Output
+      Fluent::Plugin.register_output('azure_logs_ingestion', self)
+
+      config_param :endpoint, :string
+      config_param :dcr_immutable_id, :string
+      config_param :stream_name, :string
+      config_param :gzip, :bool, default: false
+      config_param :use_msi, :bool, default: false
+      config_param :tenant_id, :string, default: ENV['AZURE_TENANT_ID']
+      config_param :client_id, :string, default: ENV['AZURE_CLIENT_ID']
+      config_param :client_secret, :string, secret: true, default: ENV['AZURE_CLIENT_SECRET']
+      config_param :authority_host, :string, default: 'https://login.microsoftonline.com'
+      config_param :logs_ingestion_scope, :string, default: 'https://monitor.azure.com/.default'
+      config_param :token_refresh_skew, :time, default: 300
+
+      config_section :buffer do
+        config_set_default :@type, 'file'
+        config_set_default :chunk_limit_size, 900 * 1024
+      end
+
+      def configure(conf)
+        super
+
+        validate_urls!
+
+        return if @use_msi
+        return if @tenant_id && @client_id && @client_secret
+
+        raise Fluent::ConfigError, 'tenant_id, client_id, and client_secret are required when use_msi is false'
+      end
+
+      def start
+        super
+        @auth = AzureLogsIngestion::Auth.new(
+          use_msi: @use_msi,
+          tenant_id: @tenant_id,
+          client_id: @client_id,
+          client_secret: @client_secret,
+          authority_host: @authority_host,
+          logs_ingestion_scope: @logs_ingestion_scope,
+          token_refresh_skew: @token_refresh_skew,
+          logger: log
+        )
+        @client = AzureLogsIngestion::Client.new(
+          endpoint: @endpoint,
+          dcr_immutable_id: @dcr_immutable_id,
+          stream_name: @stream_name,
+          logger: log
+        )
+      end
+
+      def write(chunk)
+        log.debug('building logs ingestion payload', chunk_id: dump_unique_id_hex(chunk.unique_id), gzip: @gzip)
+        payload = AzureLogsIngestion::PayloadBuilder.new(gzip: @gzip).build(chunk)
+
+        log.debug(
+          'built logs ingestion payload',
+          chunk_id: dump_unique_id_hex(chunk.unique_id),
+          record_count: payload.record_count,
+          raw_size: payload.raw_size,
+          gzip_size: payload.gzip_size,
+          content_length: payload.content_length
+        )
+
+        token = @auth.token
+        @client.send_payload(payload: payload, bearer_token: token)
+        log.debug('logs ingestion request completed', chunk_id: dump_unique_id_hex(chunk.unique_id))
+      ensure
+        payload&.close!
+      end
+
+      private
+
+      def validate_urls!
+        validate_url!(@endpoint, 'endpoint')
+        validate_url!(@authority_host, 'authority_host')
+      end
+
+      def validate_url!(value, field)
+        uri = URI.parse(value)
+        return if uri.is_a?(URI::HTTP) && uri.host
+
+        raise Fluent::ConfigError, "#{field} must be a valid HTTP or HTTPS URL"
+      rescue URI::InvalidURIError
+        raise Fluent::ConfigError, "#{field} must be a valid HTTP or HTTPS URL"
+      end
+    end
+  end
+end

data/test/helper.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+require 'test/unit'
+require 'fluent/test'
+
+Fluent::Test.setup

data/test/support/fake_azure_server.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require 'socket'
+require 'thread'
+
+class FakeAzureServer
+  Request = Struct.new(:method, :path, :headers, :body, keyword_init: true)
+
+  attr_reader :requests
+
+  def initialize(&handler)
+    @handler = handler || proc { |_request| [200, {}, ''] }
+    @requests = []
+    @mutex = Mutex.new
+    @closed = false
+  end
+
+  def start
+    @server = TCPServer.new('127.0.0.1', 0)
+    @port = @server.addr[1]
+    @thread = Thread.new { run }
+    self
+  end
+
+  def url
+    "http://127.0.0.1:#{@port}"
+  end
+
+  def stop
+    @closed = true
+    @server&.close
+    @thread&.join(1)
+  rescue IOError, Errno::EBADF
+    nil
+  end
+
+  private
+
+  def run
+    until @closed
+      begin
+        socket = @server.accept
+        handle_client(socket)
+      rescue IOError, Errno::EBADF
+        break
+      end
+    end
+  end
+
+  def handle_client(socket)
+    request_line = socket.gets("\r\n")
+    return unless request_line
+
+    method, path, = request_line.strip.split(' ', 3)
+    headers = {}
+    while (line = socket.gets("\r\n"))
+      break if line == "\r\n"
+
+      key, value = line.split(':', 2)
+      headers[key.downcase] = value.strip
+    end
+
+    body = read_body(socket, headers)
+    request = Request.new(method: method, path: path, headers: headers, body: body)
+    @mutex.synchronize { @requests << request }
+    status, response_headers, response_body = @handler.call(request)
+    write_response(socket, status, response_headers || {}, response_body || '')
+  ensure
+    socket&.close
+  end
+
+  def read_body(socket, headers)
+    length = headers.fetch('content-length', '0').to_i
+    return '' if length <= 0
+
+    socket.read(length)
+  end
+
+  def write_response(socket, status, headers, body)
+    reason = {
+      200 => 'OK',
+      400 => 'Bad Request',
+      401 => 'Unauthorized',
+      403 => 'Forbidden',
+      413 => 'Payload Too Large',
+      429 => 'Too Many Requests',
+      500 => 'Internal Server Error'
+    }.fetch(status, 'OK')
+
+    response = +"HTTP/1.1 #{status} #{reason}\r\n"
+    response << "Content-Length: #{body.bytesize}\r\n"
+    response << "Connection: close\r\n"
+    headers.each do |key, value|
+      response << "#{key}: #{value}\r\n"
+    end
+    response << "\r\n"
+    response << body
+    socket.write(response)
+  end
+end

data/test/support/helpers.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'time'
+require 'uri'
+
+class TestLogger
+  def debug(*) = nil
+  def info(*) = nil
+  def warn(*) = nil
+  def error(*) = nil
+end
+
+FakeChunk = Struct.new(:events, :chunk_id) do
+  def initialize(events, chunk_id = '0123456789ab')
+    super(events, chunk_id)
+  end
+
+  def each(&block)
+    events.each(&block)
+  end
+
+  def unique_id
+    chunk_id
+  end
+end
+
+module TestEnvHelper
+  def with_env(values)
+    previous = {}
+    values.each do |key, value|
+      previous[key] = ENV[key]
+      value.nil? ? ENV.delete(key) : ENV[key] = value
+    end
+    yield
+  ensure
+    previous.each do |key, value|
+      value.nil? ? ENV.delete(key) : ENV[key] = value
+    end
+  end
+end

data/test/test_auth_msi.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require_relative 'helper'
+require_relative 'support/fake_azure_server'
+require_relative 'support/helpers'
+require 'fluent/plugin/azure_logs_ingestion/auth'
+
+class AuthManagedIdentityTest < Test::Unit::TestCase
+  include TestEnvHelper
+
+  test 'uses app service managed identity endpoint when environment is present' do
+    server = FakeAzureServer.new do |_request|
+      [200, { 'Content-Type' => 'application/json' }, { access_token: 'msi-token', expires_on: (Time.now.to_i + 3600).to_s }.to_json]
+    end.start
+
+    with_env(
+      'IDENTITY_ENDPOINT' => "#{server.url}/msi/token",
+      'IDENTITY_HEADER' => 'identity-header-value',
+      'AZURE_LOGS_INGESTION_IMDS_ENDPOINT' => nil
+    ) do
+      auth = Fluent::Plugin::AzureLogsIngestion::Auth.new(
+        use_msi: true,
+        tenant_id: nil,
+        client_id: 'user-assigned-client-id',
+        client_secret: nil,
+        authority_host: 'https://login.microsoftonline.com',
+        logs_ingestion_scope: 'https://monitor.azure.com/.default',
+        token_refresh_skew: 300,
+        logger: TestLogger.new
+      )
+
+      assert_equal 'msi-token', auth.token
+    end
+
+    request = server.requests.first
+    query = URI.decode_www_form(URI(request.path).query).to_h
+    assert_equal 'identity-header-value', request.headers['x-identity-header']
+    assert_equal 'https://monitor.azure.com/', query['resource']
+    assert_equal 'user-assigned-client-id', query['client_id']
+  ensure
+    server&.stop
+  end
+
+  test 'uses IMDS endpoint when app service environment is absent' do
+    server = FakeAzureServer.new do |_request|
+      [200, { 'Content-Type' => 'application/json' }, { access_token: 'imds-token', expires_on: (Time.now.to_i + 3600).to_s }.to_json]
+    end.start
+
+    with_env(
+      'IDENTITY_ENDPOINT' => nil,
+      'IDENTITY_HEADER' => nil,
+      'AZURE_LOGS_INGESTION_IMDS_ENDPOINT' => "#{server.url}/metadata/identity/oauth2/token"
+    ) do
+      auth = Fluent::Plugin::AzureLogsIngestion::Auth.new(
+        use_msi: true,
+        tenant_id: nil,
+        client_id: 'user-assigned-client-id',
+        client_secret: nil,
+        authority_host: 'https://login.microsoftonline.com',
+        logs_ingestion_scope: 'https://monitor.azure.com/.default',
+        token_refresh_skew: 300,
+        logger: TestLogger.new
+      )
+
+      assert_equal 'imds-token', auth.token
+    end
+
+    request = server.requests.first
+    query = URI.decode_www_form(URI(request.path).query).to_h
+    assert_equal 'true', request.headers['metadata']
+    assert_equal 'https://monitor.azure.com/', query['resource']
+    assert_equal 'user-assigned-client-id', query['client_id']
+  ensure
+    server&.stop
+  end
+end

data/test/test_auth_service_principal.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative 'helper'
+require_relative 'support/fake_azure_server'
+require_relative 'support/helpers'
+require 'fluent/plugin/azure_logs_ingestion/auth'
+
+class AuthServicePrincipalTest < Test::Unit::TestCase
+  test 'caches service principal token until refresh is needed' do
+    server = FakeAzureServer.new do |_request|
+      [200, { 'Content-Type' => 'application/json' }, { access_token: 'token-1', expires_in: '3600' }.to_json]
+    end.start
+
+    auth = Fluent::Plugin::AzureLogsIngestion::Auth.new(
+      use_msi: false,
+      tenant_id: 'tenant-id',
+      client_id: 'client-id',
+      client_secret: 'secret',
+      authority_host: server.url,
+      logs_ingestion_scope: 'https://monitor.azure.com/.default',
+      token_refresh_skew: 300,
+      logger: TestLogger.new
+    )
+
+    assert_equal 'token-1', auth.token
+    assert_equal 'token-1', auth.token
+    assert_equal 1, server.requests.size
+    assert_match %r{/tenant-id/oauth2/v2.0/token}, server.requests.first.path
+    assert_match(/scope=https%3A%2F%2Fmonitor\.azure\.com%2F\.default/, server.requests.first.body)
+  ensure
+    server&.stop
+  end
+
+  test 'refreshes service principal token when it is already expired' do
+    tokens = %w[token-1 token-2]
+    server = FakeAzureServer.new do |_request|
+      [200, { 'Content-Type' => 'application/json' }, { access_token: tokens.shift, expires_in: '0' }.to_json]
+    end.start
+
+    auth = Fluent::Plugin::AzureLogsIngestion::Auth.new(
+      use_msi: false,
+      tenant_id: 'tenant-id',
+      client_id: 'client-id',
+      client_secret: 'secret',
+      authority_host: server.url,
+      logs_ingestion_scope: 'https://monitor.azure.com/.default',
+      token_refresh_skew: 0,
+      logger: TestLogger.new
+    )
+
+    assert_equal 'token-1', auth.token
+    assert_equal 'token-2', auth.token
+    assert_equal 2, server.requests.size
+  ensure
+    server&.stop
+  end
+end

data/test/test_out_azure_logs_ingestion.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require_relative 'helper'
+require 'fluent/test/driver/output'
+require 'fluent/plugin/out_azure_logs_ingestion'
+
+class AzureLogsIngestionOutputTest < Test::Unit::TestCase
+  BASE_CONFIG = <<~CONFIG
+    endpoint https://example.eastus-1.ingest.monitor.azure.com
+    dcr_immutable_id dcr-000a00a000a00000a000000aa000a0aa
+    stream_name Custom-MyTable
+    tenant_id test-tenant
+    client_id test-client
+    client_secret test-secret
+    <buffer>
+      @type memory
+    </buffer>
+  CONFIG
+
+  def create_driver(conf = BASE_CONFIG)
+    Fluent::Test::Driver::Output.new(Fluent::Plugin::AzureLogsIngestionOutput).configure(conf)
+  end
+
+  test 'configures minimal service principal settings' do
+    driver = create_driver
+
+    assert_equal 'https://example.eastus-1.ingest.monitor.azure.com', driver.instance.endpoint
+    assert_equal 'dcr-000a00a000a00000a000000aa000a0aa', driver.instance.dcr_immutable_id
+    assert_equal 'Custom-MyTable', driver.instance.stream_name
+  end
+
+  test 'allows managed identity without service principal secret' do
+    driver = create_driver(<<~CONFIG)
+      endpoint https://example.eastus-1.ingest.monitor.azure.com
+      dcr_immutable_id dcr-000a00a000a00000a000000aa000a0aa
+      stream_name Custom-MyTable
+      use_msi true
+      <buffer>
+        @type memory
+      </buffer>
+    CONFIG
+
+    assert_equal true, driver.instance.use_msi
+  end
+
+  test 'rejects missing credentials when use_msi is false' do
+    error = assert_raise(Fluent::ConfigError) do
+      create_driver(<<~CONFIG)
+        endpoint https://example.eastus-1.ingest.monitor.azure.com
+        dcr_immutable_id dcr-000a00a000a00000a000000aa000a0aa
+        stream_name Custom-MyTable
+        <buffer>
+          @type memory
+        </buffer>
+      CONFIG
+    end
+
+    assert_match(/tenant_id, client_id, and client_secret/, error.message)
+  end
+end