feat(aws-lambda) add support for forwarded request method/uri/body

New optional configuration properties: - `forward_request_method` - `forward_request_uri` - `forward_request_headers` - `forward_request_body` If specified, the request body sent to the invoked Lambda will contain the desired attributes of the client's request as a JSON payload. If none of the above values are specified, the upstream body contains a JSON payload which is merged from the request's body data and its query string. This is to preserve backwards compatibility with previous versions of this plugin. Original patch from Andrei Kishkevich <[email protected]> Reworked by Thibault Charbonnier <[email protected]> Signed-off-by: Thibault Charbonnier <[email protected]>
Kong · Aug 21, 2017 · df757ba · df757ba
1 parent 18771ce
commit df757ba
Show file tree

Hide file tree

Showing 4 changed files with 259 additions and 20 deletions.
diff --git a/kong/plugins/aws-lambda/handler.lua b/kong/plugins/aws-lambda/handler.lua
@@ -8,8 +8,20 @@ local http = require "resty.http"
 local cjson = require "cjson.safe"
 local public_utils = require "kong.tools.public"
 
-local ngx_req_read_body = ngx.req.read_body
+local tostring             = tostring
+local ngx_req_read_body    = ngx.req.read_body
 local ngx_req_get_uri_args = ngx.req.get_uri_args
+local ngx_req_get_headers  = ngx.req.get_headers
+local ngx_encode_base64    = ngx.encode_base64
+
+local new_tab
+do
+  local ok
+  ok, new_tab = pcall(require, "table.new")
+  if not ok then
+    new_tab = function(narr, nrec) return {} end
+  end
+end
 
 local AWS_PORT = 443
 
@@ -19,20 +31,62 @@ function AWSLambdaHandler:new()
   AWSLambdaHandler.super.new(self, "aws-lambda")
 end
 
-local function retrieve_parameters()
-  ngx_req_read_body()
-
-  return utils.table_merge(ngx_req_get_uri_args(), public_utils.get_body_args())
-end
-
 function AWSLambdaHandler:access(conf)
   AWSLambdaHandler.super.access(self)
 
-  local bodyJson = cjson.encode(retrieve_parameters())
+  local upstream_body = new_tab(0, 6)
+
+  do
+    ngx_req_read_body()
+
+    local body_args, err_code, body_raw = public_utils.get_body_info()
+    if err_code == public_utils.req_body_errors.unknown_ct then
+      -- don't know what this body MIME type is, base64 it just in case
+      body_raw = ngx_encode_base64(body_raw)
+    end
+
+    if conf.forward_request_body   or conf.forward_request_headers
+       or conf.forward_request_method or conf.forward_request_uri
+    then
+      -- new behavior to forward request method, body and its args
+      local var = ngx.var
+
+      if conf.forward_request_method then
+        upstream_body.request_method = var.request_method
+      end
+
+      if conf.forward_request_headers then
+        upstream_body.request_headers = ngx_req_get_headers()
+      end
+
+      if conf.forward_request_uri then
+        upstream_body.request_uri      = var.request_uri
+        upstream_body.request_uri_args = ngx_req_get_uri_args()
+      end
+
+      if conf.forward_request_body then
+        upstream_body.request_body      = body_raw
+        upstream_body.request_body_args = body_args
+      end
+
+    else
+      -- backwards compatible upstream body for configurations not specifying
+      -- `forward_request_*` values
+      upstream_body = utils.table_merge(ngx_req_get_uri_args(), body_args)
+    end
+  end
+
+  local upstream_body_json, err = cjson.encode(upstream_body)
+  if not upstream_body_json then
+    ngx.log(ngx.ERR, "[aws-lambda] could not JSON encode upstream body",
+                     "to forward request values: ", err)
+  end
 
   local host = string.format("lambda.%s.amazonaws.com", conf.aws_region)
   local path = string.format("/2015-03-31/functions/%s/invocations",
                             conf.function_name)
+  local port = conf.port or AWS_PORT
+
   local opts = {
     region = conf.aws_region,
     service = "lambda",
@@ -42,10 +96,12 @@ function AWSLambdaHandler:access(conf)
       ["X-Amz-Invocation-Type"] = conf.invocation_type,
       ["X-Amx-Log-Type"] = conf.log_type,
       ["Content-Type"] = "application/x-amz-json-1.1",
-      ["Content-Length"] = tostring(#bodyJson)
+      ["Content-Length"] = upstream_body_json and tostring(#upstream_body_json),
     },
-    body = bodyJson,
+    body = upstream_body_json,
     path = path,
+    host = host,
+    port = port,
     access_key = conf.aws_key,
     secret_key = conf.aws_secret,
     query = conf.qualifier and "Qualifier=" .. conf.qualifier
@@ -58,8 +114,8 @@ function AWSLambdaHandler:access(conf)
 
   -- Trigger request
   local client = http.new()
-  client:connect(host, conf.port or AWS_PORT)
   client:set_timeout(conf.timeout)
+  client:connect(host, port)
   local ok, err = client:ssl_handshake()
   if not ok then
     return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)

diff --git a/kong/plugins/aws-lambda/schema.lua b/kong/plugins/aws-lambda/schema.lua
@@ -80,5 +80,21 @@ return {
       type = "number",
       func = check_status,
     },
+    forward_request_method = {
+      type = "boolean",
+      default = false,
+    },
+    forward_request_uri = {
+      type = "boolean",
+      default = false,
+    },
+    forward_request_headers = {
+      type = "boolean",
+      default = false,
+    },
+    forward_request_body = {
+      type = "boolean",
+      default = false,
+    },
   },
 }
diff --git a/spec/03-plugins/23-aws-lambda/01-access_spec.lua b/spec/03-plugins/23-aws-lambda/01-access_spec.lua
@@ -54,6 +54,18 @@ describe("Plugin: AWS Lambda (access)", function()
       upstream_url = "http://httpbin.org"
     })
 
+    local api9 = assert(helpers.dao.apis:insert {
+      name = "lambda9.com",
+      hosts = { "lambda9.com" },
+      upstream_url = "http://httpbin.org"
+    })
+
+    local api10 = assert(helpers.dao.apis:insert {
+      name = "lambda10.com",
+      hosts = { "lambda10.com" },
+      upstream_url = "http://httpbin.org"
+    })
+
     assert(helpers.dao.plugins:insert {
       name = "aws-lambda",
       api_id = api1.id,
@@ -155,6 +167,37 @@ describe("Plugin: AWS Lambda (access)", function()
         unhandled_status = 412,
       }
     })
+    assert(helpers.dao.plugins:insert {
+      name = "aws-lambda",
+      api_id = api9.id,
+      config = {
+        port = 10001,
+        aws_key = "mock-key",
+        aws_secret = "mock-secret",
+        aws_region = "us-east-1",
+        function_name = "kongLambdaTest",
+        forward_request_method = true,
+        forward_request_uri = true,
+        forward_request_headers = true,
+        forward_request_body = true,
+      }
+    })
+
+    assert(helpers.dao.plugins:insert {
+      name = "aws-lambda",
+      api_id = api10.id,
+      config = {
+        port = 10001,
+        aws_key = "mock-key",
+        aws_secret = "mock-secret",
+        aws_region = "us-east-1",
+        function_name = "kongLambdaTest",
+        forward_request_method = true,
+        forward_request_uri = false,
+        forward_request_headers = true,
+        forward_request_body = true,
+      }
+    })
 
     assert(helpers.start_kong{
       nginx_conf = "spec/fixtures/custom_nginx.template",
@@ -183,9 +226,10 @@ describe("Plugin: AWS Lambda (access)", function()
         ["Host"] = "lambda.com"
       }
     })
-    local body = assert.res_status(200, res)
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
     assert.is_string(res.headers["x-amzn-RequestId"])
-    assert.equal([["some_value1"]], body)
+    assert.equal("some_value1", body.key1)
     assert.is_nil(res.headers["X-Amz-Function-Error"])
   end)
   it("invokes a Lambda function with POST params", function()
@@ -202,9 +246,10 @@ describe("Plugin: AWS Lambda (access)", function()
         key3 = "some_value_post3"
       }
     })
-    local body = assert.res_status(200, res)
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
     assert.is_string(res.headers["x-amzn-RequestId"])
-    assert.equal([["some_value_post1"]], body)
+    assert.equal("some_value_post1", body.key1)
   end)
   it("invokes a Lambda function with POST json body", function()
     local res = assert(client:send {
@@ -220,9 +265,10 @@ describe("Plugin: AWS Lambda (access)", function()
         key3 = "some_value_json3"
       }
     })
-    local body = assert.res_status(200, res)
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
     assert.is_string(res.headers["x-amzn-RequestId"])
-    assert.equal([["some_value_json1"]], body)
+    assert.equal("some_value_json1", body.key1)
   end)
   it("invokes a Lambda function with POST and both querystring and body params", function()
     local res = assert(client:send {
@@ -237,9 +283,130 @@ describe("Plugin: AWS Lambda (access)", function()
         key3 = "some_value_post3"
       }
     })
-    local body = assert.res_status(200, res)
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
     assert.is_string(res.headers["x-amzn-RequestId"])
-    assert.equal([["from_querystring"]], body)
+    assert.equal("from_querystring", body.key1)
+  end)
+  it("invokes a Lambda function with POST and xml payload, custom header and query partameter", function()
+    local res = assert(client:send {
+      method = "POST",
+      path = "/post?key1=from_querystring",
+      headers = {
+        ["Host"] = "lambda9.com",
+        ["Content-Type"] = "application/xml",
+        ["custom-header"] = "someheader"
+      },
+      body = "<xml/>"
+    })
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
+    assert.is_string(res.headers["x-amzn-RequestId"])
+
+    -- request_method
+    assert.equal("POST", body.request_method)
+
+    -- request_uri
+    assert.equal("/post?key1=from_querystring", body.request_uri)
+    assert.is_table(body.request_uri_args)
+
+    -- request_headers
+    assert.equal("someheader", body.request_headers["custom-header"])
+    assert.equal("lambda9.com", body.request_headers.host)
+
+    -- request_body
+    assert.equal("<xml/>", body.request_body)
+    assert.is_table(body.request_body_args)
+  end)
+  it("invokes a Lambda function with POST and json payload, custom header and query partameter", function()
+    local res = assert(client:send {
+      method = "POST",
+      path = "/post?key1=from_querystring",
+      headers = {
+        ["Host"] = "lambda10.com",
+        ["Content-Type"] = "application/json",
+        ["custom-header"] = "someheader"
+      },
+      body = { key2 = "some_value" }
+    })
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
+    assert.is_string(res.headers["x-amzn-RequestId"])
+
+    -- request_method
+    assert.equal("POST", body.request_method)
+
+    -- no request_uri
+    assert.is_nil(body.request_uri)
+    assert.is_nil(body.request_uri_args)
+
+    -- request_headers
+    assert.equal("lambda10.com", body.request_headers.host)
+    assert.equal("someheader", body.request_headers["custom-header"])
+
+    -- request_body
+    assert.equal("some_value", body.request_body_args.key2)
+    assert.is_table(body.request_body_args)
+  end)
+  it("invokes a Lambda function with POST and txt payload, custom header and query partameter", function()
+    local res = assert(client:send {
+      method = "POST",
+      path = "/post?key1=from_querystring",
+      headers = {
+        ["Host"] = "lambda9.com",
+        ["Content-Type"] = "text/plain",
+        ["custom-header"] = "someheader"
+      },
+      body = "some text"
+    })
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
+    assert.is_string(res.headers["x-amzn-RequestId"])
+
+    -- request_method
+    assert.equal("POST", body.request_method)
+
+    -- request_uri
+    assert.equal("/post?key1=from_querystring", body.request_uri)
+    assert.is_table(body.request_uri_args)
+
+    -- request_headers
+    assert.equal("someheader", body.request_headers["custom-header"])
+    assert.equal("lambda9.com", body.request_headers.host)
+
+    -- request_body
+    assert.equal("some text", body.request_body)
+    assert.is_table(body.request_body_args)
+  end)
+  it("invokes a Lambda function with POST and binary payload and custom header", function()
+    local res = assert(client:send {
+      method = "POST",
+      path = "/post?key1=from_querystring",
+      headers = {
+        ["Host"] = "lambda9.com",
+        ["Content-Type"] = "application/octet-stream",
+        ["custom-header"] = "someheader"
+      },
+      body = "01234"
+    })
+    assert.res_status(200, res)
+    local body = assert.response(res).has.jsonbody()
+    assert.is_string(res.headers["x-amzn-RequestId"])
+
+    -- request_method
+    assert.equal("POST", body.request_method)
+
+    -- request_uri
+    assert.equal("/post?key1=from_querystring", body.request_uri)
+    assert.is_table(body.request_uri_args)
+
+    -- request_headers
+    assert.equal("lambda9.com", body.request_headers.host)
+    assert.equal("someheader", body.request_headers["custom-header"])
+
+    -- request_body
+    assert.equal(ngx.encode_base64('01234'), body.request_body)
+    assert.is_table(body.request_body_args)
   end)
   it("invokes a Lambda function with POST params and Event invocation_type", function()
     local res = assert(client:send {

diff --git a/spec/fixtures/custom_nginx.template b/spec/fixtures/custom_nginx.template
@@ -305,7 +305,7 @@ http {
                 ngx.req.read_body()
                 local args = require("cjson").decode(ngx.req.get_body_data())
 
-                say(string.format("%q", qargs.key1 or args.key1), 200)
+                say(ngx.req.get_body_data(), 200)
             }
         }
     }