fluent-plugin-s3 1.7.2 → 1.8.5

data/lib/fluent/plugin/s3_extractor_gzip_command.rb

@@ -19,26 +19,15 @@ module Fluent::Plugin
       end
 
       def extract(io)
-        path = if io.respond_to?(:path)
-                 io.path
-               else
-                 temp = Tempfile.new("gzip-temp")
-                 temp.write(io.read)
-                 temp.close
-                 temp.path
-               end
-
-        stdout, succeeded = Open3.capture2("gzip #{@command_parameter} #{path}")
-        if succeeded
-          stdout
-        else
-          log.warn "failed to execute gzip command. Fallback to GzipReader. status = #{succeeded}"
-          begin
-            io.rewind
-            Zlib::GzipReader.wrap(io) do |gz|
-              gz.read
-            end
-          end
+        begin
+          extract_with_command("gzip #{@command_parameter}", io, "gzip-temp")
+        rescue SizeLimitError
+          raise
+        rescue => e
+          log.warn "gzip command execution failed: #{e.message}. Fallback to GzipExtractor."
+          io.rewind
+          extractor = GzipExtractor.new(log: log, decompression_size_limit: @decompression_size_limit)
+          extractor.extract(io)
         end
       end
     end

data/lib/fluent/plugin/s3_extractor_lzma2.rb

@@ -19,19 +19,11 @@ module Fluent::Plugin
       end
 
       def extract(io)
-        path = if io.respond_to?(path)
-                 io.path
-               else
-                 temp = Tempfile.new("xz-temp")
-                 temp.write(io.read)
-                 temp.close
-                 temp.path
-               end
-
-        stdout, succeeded = Open3.capture2("xz #{@command_parameter} #{path}")
-        if succeeded
-          stdout
-        else
+        begin
+          extract_with_command("xz #{@command_parameter}", io, "xz-temp")
+        rescue SizeLimitError
+          raise
+        rescue
           raise "Failed to extract #{path} with xz command."
         end
       end

data/lib/fluent/plugin/s3_extractor_lzo.rb

@@ -19,19 +19,11 @@ module Fluent::Plugin
       end
 
       def extract(io)
-        path = if io.respond_to?(path)
-                 io.path
-               else
-                 temp = Tempfile.new("lzop-temp")
-                 temp.write(io.read)
-                 temp.close
-                 temp.path
-               end
-
-        stdout, succeeded = Open3.capture2("lzop #{@command_parameter} #{path}")
-        if succeeded
-          stdout
-        else
+        begin
+          extract_with_command("lzop #{@command_parameter}", io, "lzop-temp")
+        rescue SizeLimitError
+          raise
+        rescue
           raise "Failed to extract #{path} with lzop command."
         end
       end

data/test/test_in_s3.rb

@@ -71,6 +71,32 @@ class S3InputTest < Test::Unit::TestCase
       end
     end
 
+    def test_aws_profile
+      d = create_driver(%[
+        aws_profile test_profile
+        s3_bucket test_bucket
+        buffer_type memory
+        <sqs>
+          queue_name test_queue
+          queue_owner_aws_account_id 123456789123
+        </sqs>
+      ])
+      assert_equal 'test_profile', d.instance.aws_profile
+    end
+
+    def test_aws_credential_process
+      d = create_driver(%[
+        aws_credential_process "/path/to/aws_signing_helper credential-process"
+        s3_bucket test_bucket
+        buffer_type memory
+        <sqs>
+          queue_name test_queue
+          queue_owner_aws_account_id 123456789123
+        </sqs>
+      ])
+      assert_equal '/path/to/aws_signing_helper credential-process', d.instance.aws_credential_process
+    end
+
     def test_without_sqs_section
       conf = %[
         aws_key_id test_key_id
@@ -315,7 +341,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -350,7 +376,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -385,7 +411,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -420,7 +446,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -455,7 +481,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -495,7 +521,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -543,7 +569,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -594,7 +620,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -614,6 +640,91 @@ EOS
     assert_equal(expected_records, events.map {|_tag, _time, record| record })
   end
 
+  data(
+    "limit_gzip"          => { type: "gzip",         input: "StringIO", limit: 10,  expected_error: true },
+    "limit_text"          => { type: "text",         input: "StringIO", limit: 10,  expected_error: true },
+    "limit_gzip_command1" => { type: "gzip_command", input: "StringIO", limit: 10,  expected_error: true },
+    "limit_gzip_command2" => { type: "gzip_command", input: "Tempfile", limit: 10,  expected_error: true },
+    "normal_gzip_command" => { type: "gzip_command", input: "Tempfile", limit: 100, expected_error: false },
+    )
+  def test_decompression_size_limit(data)
+    store_type = data[:type]
+    input_type = data[:input]
+    limit = data[:limit]
+    setup_mocks
+
+    config = <<~CONF
+      #{CONFIG}
+      check_apikey_on_start false
+      store_as #{store_type}
+      format none
+      decompression_size_limit #{limit}
+    CONF
+
+    d = create_driver(config)
+
+    s3_object = stub(Object.new)
+    s3_response = stub(Object.new)
+    s3_response.body {
+      content = "#{'a'*10}\n#{'b'*10}\n"
+
+      # Switching between Tempfile and StringIO to cover both branches of the
+      # `io.respond_to?(:path)` condition in `extract_with_command`.
+      # This ensures that:
+      # 1. The StringIO route correctly uses TextExtractor to create a protected temporary file.
+      # 2. The Tempfile route correctly limits the output size during Open3.popen3 execution.
+      io = (input_type == "Tempfile") ? Tempfile.new : StringIO.new
+
+      case store_type
+      when "gzip", "gzip_command"
+        io.binmode
+        Zlib::GzipWriter.wrap(io) { |gz|
+          gz.write content
+          gz.finish
+        }
+      when "text"
+        io.write content
+      end
+
+      io.rewind
+      io
+    }
+    s3_object.get { s3_response }
+    @s3_bucket.object(anything).at_least(1) { s3_object }
+
+    body = {
+      "Records" => [
+        {
+          "s3" => {
+            "object" => {
+              "key" => "test_key"
+            }
+          }
+        }
+      ]
+    }
+    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    @sqs_poller.get_messages(anything, anything) do |config, stats|
+      config.before_request.call(stats) if config.before_request
+      stats.request_count += 1
+      if stats.request_count >= 1
+        d.instance.instance_variable_set(:@running, false)
+      end
+      [message]
+    end
+    d.run
+
+    if data[:expected_error]
+      # Verify the protection mechanism: ensure SizeLimitError is logged.
+      assert_true d.logs.any? { |l| l.include?("Extracted data exceeds limit of #{limit} bytes") }
+    else
+      # Verify the normal execution path: ensure data is correctly extracted via Open3.popen3.
+      expected_records = [{ "message" => "#{'a'*10}\n" }, { "message" => "#{'b'*10}\n" }]
+      assert_equal(expected_records, d.events.map {|_tag, _time, record| record })
+      assert_false d.logs.any? { |l| l.include?("error_class") }
+    end
+  end
+
   def test_regexp_matching
     setup_mocks
     d = create_driver(CONFIG + "\ncheck_apikey_on_start false\nstore_as text\nformat none\nmatch_regexp .*_key?")
@@ -635,7 +746,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -664,7 +775,7 @@ EOS
         }
       ]
     }
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1
@@ -708,8 +819,8 @@ EOS
         }
       }
     }
-    
-    message = Struct::StubMessage.new(1, 1, Yajl.dump(body))
+
+    message = Struct::StubMessage.new(1, 1, JSON.generate(body))
     @sqs_poller.get_messages(anything, anything) do |config, stats|
       config.before_request.call(stats) if config.before_request
       stats.request_count += 1

data/test/test_out_s3.rb

@@ -2,6 +2,7 @@ require 'fluent/test'
 require 'fluent/test/helpers'
 require 'fluent/test/log'
 require 'fluent/test/driver/output'
+require 'fluent/version'
 require 'aws-sdk-s3'
 require 'fluent/plugin/out_s3'
 
@@ -68,6 +69,30 @@ class S3OutputTest < Test::Unit::TestCase
       assert_equal true, d.instance.check_object
     end
 
+    def test_aws_profile
+      d = create_driver(%[
+        aws_profile test_profile
+        s3_bucket test_bucket
+        path log
+        utc
+        buffer_type memory
+        time_slice_format %Y%m%d-%H
+      ])
+      assert_equal 'test_profile', d.instance.aws_profile
+    end
+
+    def test_aws_credential_process
+      d = create_driver(%[
+        aws_credential_process "/path/to/aws_signing_helper credential-process"
+        s3_bucket test_bucket
+        path log
+        utc
+        buffer_type memory
+        time_slice_format %Y%m%d-%H
+      ])
+      assert_equal '/path/to/aws_signing_helper credential-process', d.instance.aws_credential_process
+    end
+
     def test_s3_endpoint_with_valid_endpoint
       d = create_driver(CONFIG + 's3_endpoint riak-cs.example.com')
       assert_equal 'riak-cs.example.com', d.instance.s3_endpoint
@@ -109,6 +134,18 @@ class S3OutputTest < Test::Unit::TestCase
       assert(e.is_a?(Fluent::ConfigError))
     end
 
+    data('level default' => nil,
+         'level 1' => 1)
+    def test_configure_with_mime_type_zstd(level)
+      conf = CONFIG.clone
+      conf << "\nstore_as zstd\n"
+      conf << "\n<compress>\nlevel #{level}\n</compress>\n" if level
+      d = create_driver(conf)
+      assert_equal 'zst', d.instance.instance_variable_get(:@compressor).ext
+      assert_equal 'application/x-zst', d.instance.instance_variable_get(:@compressor).content_type
+      assert_equal (level || 3), d.instance.instance_variable_get(:@compressor).instance_variable_get(:@compress_config).level
+    end
+
     def test_configure_with_path_style
       conf = CONFIG.clone
       conf << "\nforce_path_style true\n"
@@ -456,6 +493,152 @@ EOC
     FileUtils.rm_f(s3_local_file_path)
   end
 
+  def test_write_with_zstd
+    setup_mocks(true)
+    s3_local_file_path = "/tmp/s3-test.zst"
+
+    expected_s3path = "log/events/ts=20110102-13/events_0-#{Socket.gethostname}.zst"
+
+    setup_s3_object_mocks(s3_local_file_path: s3_local_file_path, s3path: expected_s3path)
+
+    config = CONFIG_TIME_SLICE + "\nstore_as zstd\n"
+    d = create_time_sliced_driver(config)
+
+    time = event_time("2011-01-02 13:14:15 UTC")
+    d.run(default_tag: "test") do
+      d.feed(time, { "a" => 1 })
+      d.feed(time, { "a" => 2 })
+    end
+
+    File.open(s3_local_file_path, 'rb') do |file|
+      compressed_data = file.read
+      uncompressed_data = Zstd.decompress(compressed_data)
+      expected_data = %[2011-01-02T13:14:15Z\ttest\t{"a":1}\n] +
+                      %[2011-01-02T13:14:15Z\ttest\t{"a":2}\n]
+      assert_equal expected_data, uncompressed_data
+    end
+    FileUtils.rm_f(s3_local_file_path)
+  end
+
+  def test_no_compress_by_gzip_with_compressed_buffer
+    setup_mocks(true)
+    s3_local_file_path = "/tmp/s3-test.gz"
+    setup_s3_object_mocks(s3_local_file_path: s3_local_file_path)
+
+    d = create_time_sliced_driver(CONFIG_TIME_SLICE + <<~EOF)
+      <buffer tag,time>
+        @type memory
+        compress gzip
+        timekey 3600
+        timekey_use_utc true
+      </buffer>
+    EOF
+
+    # GzipCompressor should not use Zlib::GzipWriter
+    dont_allow(Zlib::GzipWriter).new
+
+    time = event_time("2011-01-02 13:14:15 UTC")
+    d.run(default_tag: "test") do
+      d.feed(time, {"a"=>1})
+      d.feed(time, {"a"=>2})
+    end
+
+    data = ""
+    File.open(s3_local_file_path, "rb") do |f|
+      until f.eof?
+        gz = Zlib::GzipReader.new(f)
+        data << gz.read
+
+        unused = gz.unused
+        gz.finish
+        f.pos -= unused.length if unused
+      end
+    end
+    assert_equal %[2011-01-02T13:14:15Z\ttest\t{"a":1}\n] +
+                 %[2011-01-02T13:14:15Z\ttest\t{"a":2}\n],
+                 data
+    FileUtils.rm_f(s3_local_file_path)
+  end
+
+  def test_no_compress_by_gzip_command_with_compressed_buffer
+    setup_mocks(true)
+    s3_local_file_path = "/tmp/s3-test.gz"
+    setup_s3_object_mocks(s3_local_file_path: s3_local_file_path)
+
+    d = create_time_sliced_driver(CONFIG_TIME_SLICE + <<~EOF)
+      store_as gzip_command
+      <buffer tag,time>
+        @type memory
+        compress gzip
+        timekey 3600
+        timekey_use_utc true
+      </buffer>
+    EOF
+
+    # GzipCommandCompressor should not use Kernel.system and Zlib::GzipWriter
+    dont_allow(Kernel).system
+    dont_allow(Zlib::GzipWriter).new
+
+    time = event_time("2011-01-02 13:14:15 UTC")
+    d.run(default_tag: "test") do
+      d.feed(time, {"a"=>1})
+      d.feed(time, {"a"=>2})
+    end
+
+    data = ""
+    File.open(s3_local_file_path, "rb") do |f|
+      until f.eof?
+        gz = Zlib::GzipReader.new(f)
+        data << gz.read
+
+        unused = gz.unused
+        gz.finish
+        f.pos -= unused.length if unused
+      end
+    end
+    assert_equal %[2011-01-02T13:14:15Z\ttest\t{"a":1}\n] +
+                 %[2011-01-02T13:14:15Z\ttest\t{"a":2}\n],
+                 data
+    FileUtils.rm_f(s3_local_file_path)
+  end
+
+  def test_no_compress_by_zstd_with_compressed_buffer
+    pend "This test require Fluentd 1.19.0 or later which supports zstd buffer compressed." if Gem::Version.new(Fluent::VERSION) < Gem::Version.new('1.19.0')
+
+    setup_mocks(true)
+    s3_local_file_path = "/tmp/s3-test.zst"
+    expected_s3path = "log/events/ts=20110102-13/events_0-#{Socket.gethostname}.zst"
+    setup_s3_object_mocks(s3_local_file_path: s3_local_file_path, s3path: expected_s3path)
+
+    d = create_time_sliced_driver(CONFIG_TIME_SLICE + <<~EOF)
+      store_as zstd
+      <buffer tag,time>
+        @type memory
+        compress zstd
+        timekey 3600
+        timekey_use_utc true
+      </buffer>
+    EOF
+
+    # ZstdCompressor should not use Zstd.compress
+    dont_allow(Zstd).compress
+
+    time = event_time("2011-01-02 13:14:15 UTC")
+    d.run(default_tag: "test") do
+      d.feed(time, {"a"=>1})
+      d.feed(time, {"a"=>2})
+    end
+
+    File.open(s3_local_file_path, 'rb') do |file|
+      compressed_data = file.read
+      uncompressed_data = Zstd.decompress(compressed_data)
+      expected_data = %[2011-01-02T13:14:15Z\ttest\t{"a":1}\n] +
+                      %[2011-01-02T13:14:15Z\ttest\t{"a":2}\n]
+      assert_equal expected_data, uncompressed_data
+    end
+    FileUtils.rm_f(s3_local_file_path)
+  end
+
   class MockResponse
     attr_reader :data
 
@@ -764,6 +947,92 @@ EOC
     assert_equal(expected_credentials, credentials)
   end
 
+  def test_web_identity_credentials_with_region_and_sts_http_proxy
+    expected_credentials = Aws::Credentials.new("test_key", "test_secret")
+    expected_region = "ap-northeast-1"
+    expected_sts_http_proxy = 'http://example.com'
+    sts_client = Aws::STS::Client.new(region: expected_region, http_proxy: expected_sts_http_proxy)
+    mock(Aws::STS::Client).new(region:expected_region, http_proxy: expected_sts_http_proxy){ sts_client }
+    mock(Aws::AssumeRoleWebIdentityCredentials).new({ role_arn: "test_arn",
+                                           role_session_name: "test_session",
+                                           web_identity_token_file: "test_file",
+                                           client: sts_client,
+                                           sts_http_proxy: expected_sts_http_proxy }){
+      expected_credentials
+    }
+    config = CONFIG_TIME_SLICE.split("\n").reject{|x| x =~ /.+aws_.+/}.join("\n")
+    config += %[
+      s3_region #{expected_region}
+      <web_identity_credentials>
+        role_arn test_arn
+        role_session_name test_session
+        web_identity_token_file test_file
+        sts_http_proxy #{expected_sts_http_proxy}
+      </web_identity_credentials>
+    ]
+    d = create_time_sliced_driver(config)
+    assert_nothing_raised { d.run {} }
+    client = d.instance.instance_variable_get(:@s3).client
+    credentials = client.config.credentials
+    assert_equal(expected_credentials, credentials)
+  end
+
+  def test_web_identity_credentials_with_sts_http_proxy
+    expected_credentials = Aws::Credentials.new("test_key", "test_secret")
+    expected_sts_http_proxy = 'http://example.com'
+    sts_client = Aws::STS::Client.new(region: "us-east-1", http_proxy: expected_sts_http_proxy)
+    mock(Aws::STS::Client).new(region: "us-east-1", http_proxy: expected_sts_http_proxy){ sts_client }
+    mock(Aws::AssumeRoleWebIdentityCredentials).new({ role_arn: "test_arn",
+                                           role_session_name: "test_session",
+                                           web_identity_token_file: "test_file",
+                                           client: sts_client,
+                                           sts_http_proxy: expected_sts_http_proxy }){
+      expected_credentials
+    }
+    config = CONFIG_TIME_SLICE.split("\n").reject{|x| x =~ /.+aws_.+/}.join("\n")
+    config += %[
+      <web_identity_credentials>
+        role_arn test_arn
+        role_session_name test_session
+        web_identity_token_file test_file
+        sts_http_proxy #{expected_sts_http_proxy}
+      </web_identity_credentials>
+    ]
+    d = create_time_sliced_driver(config)
+    assert_nothing_raised { d.run {} }
+    client = d.instance.instance_variable_get(:@s3).client
+    credentials = client.config.credentials
+    assert_equal(expected_credentials, credentials)
+  end
+
+  def test_web_identity_credentials_with_sts_endpoint_url
+    expected_credentials = Aws::Credentials.new("test_key", "test_secret")
+    expected_sts_endpoint_url = 'http://example.com'
+    sts_client = Aws::STS::Client.new(region: "us-east-1", endpoint: expected_sts_endpoint_url)
+    mock(Aws::STS::Client).new(region: "us-east-1", endpoint: expected_sts_endpoint_url){ sts_client }
+    mock(Aws::AssumeRoleWebIdentityCredentials).new({ role_arn: "test_arn",
+                                           role_session_name: "test_session",
+                                           web_identity_token_file: "test_file",
+                                           client: sts_client,
+                                           sts_endpoint_url: expected_sts_endpoint_url }){
+      expected_credentials
+    }
+    config = CONFIG_TIME_SLICE.split("\n").reject{|x| x =~ /.+aws_.+/}.join("\n")
+    config += %[
+      <web_identity_credentials>
+        role_arn test_arn
+        role_session_name test_session
+        web_identity_token_file test_file
+        sts_endpoint_url #{expected_sts_endpoint_url}
+      </web_identity_credentials>
+    ]
+    d = create_time_sliced_driver(config)
+    assert_nothing_raised { d.run {} }
+    client = d.instance.instance_variable_get(:@s3).client
+    credentials = client.config.credentials
+    assert_equal(expected_credentials, credentials)
+  end
+
   def test_web_identity_credentials_with_sts_region
     expected_credentials = Aws::Credentials.new("test_key", "test_secret")
     sts_client = Aws::STS::Client.new(region: 'us-east-1')