diff --git a/lib/aws_lambda_ric.rb b/lib/aws_lambda_ric.rb index b4e4185..d898664 100755 --- a/lib/aws_lambda_ric.rb +++ b/lib/aws_lambda_ric.rb @@ -26,8 +26,8 @@ class LambdaRunner ENV_VAR_RUNTIME_API = 'AWS_LAMBDA_RUNTIME_API' - def initialize(runtime_server_addr) - @lambda_server = LambdaServer.new(runtime_server_addr) + def initialize(runtime_server_addr, user_agent) + @lambda_server = LambdaServer.new(runtime_server_addr, user_agent) @runtime_loop_active = true # if false, we will exit the program @exit_code = 0 end diff --git a/lib/aws_lambda_ric/bootstrap.rb b/lib/aws_lambda_ric/bootstrap.rb index ffca6e3..70d0bcd 100644 --- a/lib/aws_lambda_ric/bootstrap.rb +++ b/lib/aws_lambda_ric/bootstrap.rb @@ -30,9 +30,16 @@ def self.bootstrap_handler end app_root = Dir.pwd handler = ARGV[0] - lambda_runner = AwsLambdaRuntimeInterfaceClient::LambdaRunner.new(fetch_runtime_server) + lambda_runner = AwsLambdaRuntimeInterfaceClient::LambdaRunner.new(fetch_runtime_server, get_user_agent) puts "Executing '#{handler}' in function directory '#{app_root}'" lambda_runner.run(app_root, handler) end + def self.get_user_agent + ruby_version = RUBY_VERSION.to_s + version = AwsLambdaRuntimeInterfaceClient::VERSION + + "aws-lambda-ruby/#{ruby_version}-#{version}" + end + end \ No newline at end of file diff --git a/lib/aws_lambda_ric/lambda_server.rb b/lib/aws_lambda_ric/lambda_server.rb index 7c36c64..5b3c14c 100644 --- a/lib/aws_lambda_ric/lambda_server.rb +++ b/lib/aws_lambda_ric/lambda_server.rb @@ -10,21 +10,22 @@ class LambdaServer LAMBDA_DEFAULT_SERVER_ADDRESS = '127.0.0.1:9001' LAMBDA_RUNTIME_API_VERSION = '2018-06-01' - MAX_HEADER_SIZE = 1024 * 1024 - LONG_TIMEOUT = 1_000_000 + MAX_HEADER_SIZE_BYTES = 1024 * 1024 + LONG_TIMEOUT_MS = 1_000_000 - def initialize(server_address) + def initialize(server_address, user_agent) server_address ||= LAMBDA_DEFAULT_SERVER_ADDRESS - @server_address = 'http://' + server_address + '/' + LAMBDA_RUNTIME_API_VERSION + @server_address = "http://#{server_address}/#{LAMBDA_RUNTIME_API_VERSION}" + @user_agent = user_agent end def next_invocation next_invocation_uri = URI(@server_address + '/runtime/invocation/next') begin http = Net::HTTP.new(next_invocation_uri.host, next_invocation_uri.port) - http.read_timeout = LONG_TIMEOUT + http.read_timeout = LONG_TIMEOUT_MS resp = http.start do |connection| - connection.get(next_invocation_uri.path) + connection.get(next_invocation_uri.path, { 'User-Agent' => @user_agent }) end if resp.is_a?(Net::HTTPSuccess) request_id = resp['Lambda-Runtime-Aws-Request-Id'] @@ -51,7 +52,7 @@ def send_response(request_id:, response_object:, content_type: 'application/json Net::HTTP.post( response_uri, response_object, - { 'Content-Type' => content_type } + { 'Content-Type' => content_type, 'User-Agent' => @user_agent } ) rescue StandardError => e raise LambdaErrors::LambdaRuntimeError.new(e) @@ -61,8 +62,8 @@ def send_response(request_id:, response_object:, content_type: 'application/json def send_error_response(request_id:, error_object:, error:, xray_cause:) response_uri = URI(@server_address + "/runtime/invocation/#{request_id}/error") begin - headers = { 'Lambda-Runtime-Function-Error-Type' => error.runtime_error_type } - headers['Lambda-Runtime-Function-XRay-Error-Cause'] = xray_cause if xray_cause.bytesize < MAX_HEADER_SIZE + headers = { 'Lambda-Runtime-Function-Error-Type' => error.runtime_error_type, 'User-Agent' => @user_agent } + headers['Lambda-Runtime-Function-XRay-Error-Cause'] = xray_cause if xray_cause.bytesize < MAX_HEADER_SIZE_BYTES Net::HTTP.post( response_uri, error_object.to_json, @@ -79,7 +80,7 @@ def send_init_error(error_object:, error:) Net::HTTP.post( uri, error_object.to_json, - { 'Lambda-Runtime-Function-Error-Type' => error.runtime_error_type } + { 'Lambda-Runtime-Function-Error-Type' => error.runtime_error_type, 'User-Agent' => @user_agent } ) rescue StandardError => e raise LambdaErrors::LambdaRuntimeInitError.new(e) diff --git a/test/unit/lambda_server_test.rb b/test/unit/lambda_server_test.rb index 1ba3926..2edbf43 100644 --- a/test/unit/lambda_server_test.rb +++ b/test/unit/lambda_server_test.rb @@ -11,13 +11,15 @@ def setup @request_id = 'test_id' @error = LambdaErrors::LambdaRuntimeError.new(StandardError.new('User error, replace user')) @error_uri = URI("http://#{@server_address}/2018-06-01/runtime/invocation/#{@request_id}/error") - @under_test = LambdaServer.new(@server_address) + @mock_user_agent = 'mock-user-agent' + @under_test = LambdaServer.new(@server_address, @mock_user_agent) end def test_post_invocation_error_with_large_xray_cause large_xray_cause = ('a' * 1024 * 1024)[0..-2] headers = {'Lambda-Runtime-Function-Error-Type' => @error.runtime_error_type, - 'Lambda-Runtime-Function-XRay-Error-Cause' => large_xray_cause} + 'Lambda-Runtime-Function-XRay-Error-Cause' => large_xray_cause, + 'User-Agent' => @mock_user_agent} post_mock = Minitest::Mock.new post_mock.expect :call, nil, [@error_uri, @error.to_lambda_response.to_json, headers] @@ -35,7 +37,8 @@ def test_post_invocation_error_with_large_xray_cause def test_post_invocation_error_with_too_large_xray_cause too_large_xray_cause = 'a' * 1024 * 1024 - headers = {'Lambda-Runtime-Function-Error-Type' => @error.runtime_error_type} + headers = {'Lambda-Runtime-Function-Error-Type' => @error.runtime_error_type, + 'User-Agent' => @mock_user_agent} post_mock = Minitest::Mock.new post_mock.expect :call, nil, [@error_uri, @error.to_lambda_response.to_json, headers]