Skip to content

Commit

Permalink
Merge pull request #4137 from daipom/in_tcp-add-message_length_limit
Browse files Browse the repository at this point in the history
`in_tcp`: Add `message_length_limit` to drop large incoming data
  • Loading branch information
ashie authored Apr 17, 2023
2 parents c989079 + 0d9c320 commit 11c9550
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 0 deletions.
43 changes: 43 additions & 0 deletions lib/fluent/plugin/in_tcp.rb
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ class TcpInput < Input
desc "The field name of the client's address."
config_param :source_address_key, :string, default: nil

# Setting default to nil for backward compatibility
desc "The max bytes of message."
config_param :message_length_limit, :size, default: nil

config_param :blocking_timeout, :time, default: 0.5

desc 'The payload is read up to this character.'
Expand Down Expand Up @@ -102,6 +106,7 @@ def start

log.info "listening tcp socket", bind: @bind, port: @port
del_size = @delimiter.length
discard_till_next_delimiter = false
if @_extract_enabled && @_extract_tag_key
server_create(:in_tcp_server_single_emit, @port, bind: @bind, resolve_name: !!@source_hostname_key, send_keepalive_packet: @send_keepalive_packet) do |data, conn|
unless check_client(conn)
Expand All @@ -116,6 +121,16 @@ def start
msg = buf[pos...i]
pos = i + del_size

if discard_till_next_delimiter
discard_till_next_delimiter = false
next
end

if !@message_length_limit.nil? && @message_length_limit < msg.bytesize
log.info "The received data is larger than 'message_length_limit', dropped:", limit: @message_length_limit, size: msg.bytesize, head: msg[...32]
next
end

@parser.parse(msg) do |time, record|
unless time && record
log.warn "pattern not matched", message: msg
Expand All @@ -131,6 +146,15 @@ def start
end
end
buf.slice!(0, pos) if pos > 0
# If the buffer size exceeds the limit here, it means that the next message will definitely exceed the limit.
# So we should clear the buffer here. Otherwise, it will keep storing useless data until the next delimiter comes.
if !@message_length_limit.nil? && @message_length_limit < buf.bytesize
log.info "The buffer size exceeds 'message_length_limit', cleared:", limit: @message_length_limit, size: buf.bytesize, head: buf[...32]
buf.clear
# We should discard the subsequent data until the next delimiter comes.
discard_till_next_delimiter = true
next
end
end
else
server_create(:in_tcp_server_batch_emit, @port, bind: @bind, resolve_name: !!@source_hostname_key, send_keepalive_packet: @send_keepalive_packet) do |data, conn|
Expand All @@ -147,6 +171,16 @@ def start
msg = buf[pos...i]
pos = i + del_size

if discard_till_next_delimiter
discard_till_next_delimiter = false
next
end

if !@message_length_limit.nil? && @message_length_limit < msg.bytesize
log.info "The received data is larger than 'message_length_limit', dropped:", limit: @message_length_limit, size: msg.bytesize, head: msg[...32]
next
end

@parser.parse(msg) do |time, record|
unless time && record
log.warn "pattern not matched", message: msg
Expand All @@ -161,6 +195,15 @@ def start
end
router.emit_stream(@tag, es)
buf.slice!(0, pos) if pos > 0
# If the buffer size exceeds the limit here, it means that the next message will definitely exceed the limit.
# So we should clear the buffer here. Otherwise, it will keep storing useless data until the next delimiter comes.
if !@message_length_limit.nil? && @message_length_limit < buf.bytesize
log.info "The buffer size exceeds 'message_length_limit', cleared:", limit: @message_length_limit, size: buf.bytesize, head: buf[...32]
buf.clear
# We should discard the subsequent data until the next delimiter comes.
discard_till_next_delimiter = true
next
end
end
end
end
Expand Down
72 changes: 72 additions & 0 deletions test/plugin/test_in_tcp.rb
Original file line number Diff line number Diff line change
Expand Up @@ -253,4 +253,76 @@ def create_tcp_socket(host, port, &block)
assert_equal 'hello', event[2]['msg']
end
end

sub_test_case "message_length_limit" do
data("batch_emit", { extract: "" }, keep: true)
data("single_emit", { extract: "<extract>\ntag_key tag\n</extract>\n" }, keep: true)
test "drop records exceeding limit" do |data|
message_length_limit = 10
d = create_driver(base_config + %!
message_length_limit #{message_length_limit}
<parse>
@type none
</parse>
#{data[:extract]}
!)
d.run(expect_records: 2, timeout: 10) do
create_tcp_socket('127.0.0.1', @port) do |sock|
sock.send("a" * message_length_limit + "\n", 0)
sock.send("b" * (message_length_limit + 1) + "\n", 0)
sock.send("c" * (message_length_limit - 1) + "\n", 0)
end
end

expected_records = [
"a" * message_length_limit,
"c" * (message_length_limit - 1)
]
actual_records = d.events.collect do |event|
event[2]["message"]
end

assert_equal expected_records, actual_records
end

test "clear buffer and discard the subsequent data until the next delimiter" do |data|
message_length_limit = 12
d = create_driver(base_config + %!
message_length_limit #{message_length_limit}
delimiter ";"
<parse>
@type json
</parse>
#{data[:extract]}
!)
d.run(expect_records: 1, timeout: 10) do
create_tcp_socket('127.0.0.1', @port) do |sock|
sock.send('{"message":', 0)
sock.send('"hello', 0)
sleep 1 # To make the server read data and clear the buffer here.
sock.send('world!"};', 0) # This subsequent data must be discarded so that a parsing failure doesn't occur.
sock.send('{"k":"v"};', 0) # This will succeed to parse.
end
end

logs = d.logs.collect do |log|
log.gsub(/\A\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} [-+]\d{4} /, "")
end
actual_records = d.events.collect do |event|
event[2]
end

assert_equal(
{
# Asserting that '[warn]: pattern not matched message="world!\"}"' warning does not occur.
logs: ['[info]: The buffer size exceeds \'message_length_limit\', cleared: limit=12 size=17 head="{\"message\":\"hello"' + "\n"],
records: [{"k" => "v"}],
},
{
logs: logs[1..],
records: actual_records,
}
)
end
end
end

0 comments on commit 11c9550

Please sign in to comment.