enabling lambda functions to return custom status, body, and headers

kong/plugins/aws-lambda/handler.lua

@@ -12,6 +12,7 @@ local public_utils = require "kong.tools.public"
 
 
 local tostring             = tostring
+local tonumber             = tonumber
 local pairs                = pairs
 local type                 = type
 local fmt                  = string.format
@@ -38,6 +39,53 @@ local server_header = meta._SERVER_TOKENS
 local AWS_PORT = 443
 
 
+--[[
+  Response format should be
+  {
+      "statusCode": httpStatusCode,
+      "headers": { "headerName": "headerValue", ... },
+      "body": "..."
+  }
+--]]
+local function validate_custom_response(response)
+  if type(response.statusCode) ~= "number" then
+    return nil, "statusCode must be a number"
+  end
+
+  if response.headers ~= nil and type(response.headers) ~= "table" then
+    return nil, "headers must be a table"
+  end
+
+  if response.body ~= nil and type(response.body) ~= "string" then
+    return nil, "body must be a string"
+  end
+
+  return true
+end
+
+
+local function extract_proxy_response(content)
+  local serialized_content, err = cjson.decode(content)
+  if not serialized_content then
+    return nil, err
+  end
+
+  local ok, err = validate_custom_response(serialized_content)
+  if not ok then
+    return nil, err
+  end
+
+  local headers = serialized_content.headers or {}
+  local body = serialized_content.body or ""
+  headers["Content-Length"] = #body
+
+  return {
+    status_code = tonumber(serialized_content.statusCode),
+    body = body,
+    headers = headers,
+  }
+end
+
 local function send(status, content, headers)
   ngx.status = status
 
@@ -187,13 +235,27 @@ function AWSLambdaHandler:access(conf)
   end
 
   local status
-  if conf.unhandled_status
-     and headers["X-Amz-Function-Error"] == "Unhandled"
-  then
-    status = conf.unhandled_status
 
-  else
-    status = res.status
+  if conf.is_proxy_integration then
+    local proxy_response, err = extract_proxy_response(content)
+    if not proxy_response then
+      return responses.send_HTTP_BAD_GATEWAY(err)
+    end
+
+    headers = utils.table_merge(headers, proxy_response.headers)
+    content = proxy_response.body
+    status = proxy_response.status_code
+  end
+
+  if not status then
+    if conf.unhandled_status
+      and headers["X-Amz-Function-Error"] == "Unhandled"
+    then
+      status = conf.unhandled_status
+
+    else
+      status = res.status
+    end
   end
 
   local ctx = ngx.ctx

kong/plugins/aws-lambda/schema.lua

@@ -97,5 +97,9 @@ return {
       type = "boolean",
       default = false,
     },
+    is_proxy_integration = {
+      type = "boolean",
+      default = false,
+    },
   },
 }

kong/tools/responses.lua

@@ -59,6 +59,7 @@ local _M = {
     HTTP_CONFLICT = 409,
     HTTP_UNSUPPORTED_MEDIA_TYPE = 415,
     HTTP_INTERNAL_SERVER_ERROR = 500,
+    HTTP_BAD_GATEWAY = 502,
     HTTP_SERVICE_UNAVAILABLE = 503,
   }
 }
@@ -73,6 +74,7 @@ local _M = {
 -- @field status_codes.HTTP_INTERNAL_SERVER_ERROR Always "Internal Server Error"
 -- @field status_codes.HTTP_METHOD_NOT_ALLOWED Always "Method not allowed"
 -- @field status_codes.HTTP_SERVICE_UNAVAILABLE Default: "Service unavailable"
+-- @field status_codes.HTTP_BAD_GATEWAY Always: "Bad Gateway"
 local response_default_content = {
   [_M.status_codes.HTTP_UNAUTHORIZED] = function(content)
     return content or "Unauthorized"
@@ -92,6 +94,9 @@ local response_default_content = {
   [_M.status_codes.HTTP_SERVICE_UNAVAILABLE] = function(content)
     return content or "Service unavailable"
   end,
+  [_M.status_codes.HTTP_BAD_GATEWAY] = function(content)
+    return "Bad Gateway"
+  end,
 }
 
 -- Return a closure which will be usable to respond with a certain status code.
@@ -116,7 +121,7 @@ local function send_response(status_code)
       coroutine.yield()
     end
 
-    if status_code == _M.status_codes.HTTP_INTERNAL_SERVER_ERROR then
+    if status_code == _M.status_codes.HTTP_INTERNAL_SERVER_ERROR or status_code == _M.status_codes.HTTP_BAD_GATEWAY then
       if content then
         ngx.log(ngx.ERR, tostring(content))
       end

spec/01-unit/009-responses_spec.lua

@@ -59,7 +59,7 @@ describe("Response helpers", function()
       end)
     end
   end)
-  it("calls `ngx.log` if and only if a 500 status code was given", function()
+  it("calls `ngx.log` if and only if a 500 or 502 status code was given", function()
     responses.send_HTTP_BAD_REQUEST()
     assert.stub(ngx.log).was_not_called()
 
@@ -69,8 +69,14 @@ describe("Response helpers", function()
     responses.send_HTTP_INTERNAL_SERVER_ERROR()
     assert.stub(ngx.log).was_not_called()
 
+    responses.send_HTTP_BAD_GATEWAY()
+    assert.stub(ngx.log).was_not_called()
+
     responses.send_HTTP_INTERNAL_SERVER_ERROR("error")
     assert.stub(ngx.log).was_called()
+
+    responses.send_HTTP_BAD_GATEWAY("error")
+    assert.stub(ngx.log).was_called(2)
   end)
 
   it("don't call `ngx.log` if a 503 status code was given", function()

spec/03-plugins/23-aws-lambda/01-access_spec.lua

@@ -1,3 +1,4 @@
+local cjson = require "cjson"
 local helpers = require "spec.helpers"
 local meta    = require "kong.meta"
 
@@ -67,6 +68,30 @@ for _, strategy in helpers.each_strategy() do
         service     = service10
       }
 
+      local service11 = bp.services:insert({
+        protocol = "http",
+        host     = "httpbin.org",
+        port     = 80,
+      })
+
+      local route11 = bp.routes:insert {
+        hosts       = { "lambda11.com" },
+        protocols   = { "http", "https" },
+        service     = service11
+      }
+
+      local service12 = bp.services:insert({
+        protocol = "http",
+        host     = "httpbin.org",
+        port     = 80,
+      })
+
+      local route12 = bp.routes:insert {
+        hosts       = { "lambda12.com" },
+        protocols   = { "http", "https" },
+        service     = service12
+      }
+
       bp.plugins:insert {
         name     = "aws-lambda",
         route_id = route1.id,
@@ -201,6 +226,32 @@ for _, strategy in helpers.each_strategy() do
         }
       }
 
+      bp.plugins:insert {
+        name     = "aws-lambda",
+        route_id = route11.id,
+        config                    = {
+          port                    = 10001,
+          aws_key                 = "mock-key",
+          aws_secret              = "mock-secret",
+          aws_region              = "us-east-1",
+          function_name           = "kongLambdaTest",
+          is_proxy_integration    = true
+        }
+      }
+
+      bp.plugins:insert {
+        name     = "aws-lambda",
+        route_id = route12.id,
+        config                    = {
+          port                    = 10001,
+          aws_key                 = "mock-key",
+          aws_secret              = "mock-secret",
+          aws_region              = "us-east-1",
+          function_name           = "functionWithBadJSON",
+          is_proxy_integration    = true
+        }
+      }
+
       assert(helpers.start_kong{
         database   = strategy,
         nginx_conf = "spec/fixtures/custom_nginx.template",
@@ -527,5 +578,107 @@ for _, strategy in helpers.each_strategy() do
       assert.equal(65, tonumber(res.headers["Content-Length"]))
     end)
 
+    it("returns the proper status code when setting custom response body from Lambda", function()
+      -- the lambda function must return a string
+      -- for the custom response "body" property
+      local body = cjson.encode({
+        key1 = "some_value_post1",
+        key2 = "some_value_post2",
+        key3 = "some_value_post3"
+      })
+      local res = assert(proxy_client:send {
+        method  = "POST",
+        path    = "/post",
+        headers = {
+          ["Host"]         = "lambda11.com",
+          ["Content-Type"] = "application/json"
+        },
+        body = {
+          statusCode = 201,
+          body = body,
+          headers = {
+            ["X-Custom-Header"] = "Hello world!"
+          }
+        }
+      })
+      local res_body = assert.res_status(201, res)
+      assert.equal(79, tonumber(res.headers["Content-Length"]))
+      assert.equal("Hello world!", res.headers["X-Custom-Header"])
+      assert.equal(body, res_body)
+    end)
+
+    it("returns proper status code with no body or headers", function()
+      local res = assert(proxy_client:send {
+        method  = "POST",
+        path    = "/post",
+        headers = {
+          ["Host"]         = "lambda11.com",
+          ["Content-Type"] = "application/json"
+        },
+        body = {
+          statusCode = 201,
+        }
+      })
+      local body = assert.res_status(201, res)
+      assert.equal(0, tonumber(res.headers["Content-Length"]))
+      assert.equal(nil, res.headers["X-Custom-Header"])
+      assert.equal("", body)
+    end)
+
+    it("override duplicate headers with value from the custom response", function()
+      -- the default "x-amzn-RequestId" returned is "foo"
+      -- let's check it is overriden with a custom value
+      local headers = {
+        ["x-amzn-RequestId"] = "bar"
+      }
+      local res = assert(proxy_client:send {
+        method  = "POST",
+        path    = "/post",
+        headers = {
+          ["Host"]         = "lambda11.com",
+          ["Content-Type"] = "application/json"
+        },
+        body = {
+          statusCode = 201,
+          headers = headers,
+        }
+      })
+
+      assert.res_status(201, res)
+      assert.equal("bar", res.headers["x-amzn-RequestId"])
+    end)
+
+    it("errors with 502 with when 'body' on custom response is not a string", function()
+      local res = assert(proxy_client:send {
+        method  = "POST",
+        path    = "/post",
+        headers = {
+          ["Host"]         = "lambda11.com",
+          ["Content-Type"] = "application/json"
+        },
+        body = {
+          statusCode = 201,
+          body = 1234,
+        }
+      })
+
+      assert.res_status(502, res)
+      local b = assert.response(res).has.jsonbody()
+      assert.equal("Bad Gateway", b.message)
+    end)
+
+    it("errors with 502 with when response from lambda is not json", function()
+      local res = assert(proxy_client:send {
+        method  = "POST",
+        path    = "/post",
+        headers = {
+          ["Host"] = "lambda12.com"
+        }
+      })
+
+      assert.res_status(502, res)
+      local b = assert.response(res).has.jsonbody()
+      assert.equal("Bad Gateway", b.message)
+    end)
   end)
 end

spec/fixtures/custom_nginx.template

@@ -244,7 +244,11 @@ http {
                   end
                   ngx.status = status
 
-                  if type(res) == 'string' then
+                  if string.match(ngx.var.uri, "functionWithBadJSON") then
+                    local badRes = "{\"foo\":\"bar\""
+                    ngx.header["Content-Length"] = #badRes + 1
+                    ngx.say(badRes)
+                  elseif type(res) == 'string' then
                     ngx.header["Content-Length"] = #res + 1
                     ngx.say(res)
 
@@ -253,6 +257,7 @@ http {
                     ngx.header['Content-Length'] = 0
                   end
 
+
                   ngx.exit(0)
                 end