Added User-Agent header (#9)

Co-authored-by: Stefan Botez <[email protected]>

Author: stefanbotez

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

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

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)

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]