From 0f54bd2c88dcceaa1a8896ef8132ca0f7e62e4ef Mon Sep 17 00:00:00 2001
From: Keery Nie <windmgc@gmail.com>
Date: Fri, 2 Aug 2024 15:04:37 +0800
Subject: [PATCH] fix(request): build request body by operation payload field
 (#126)

This is a follow-up fix for #120 .

The PR modifies the request build again with the following changes: when the payload field is specified inside the input schema, the request body will be built based on the parameter that specified by the payload field.

This behavior imitates the behavior defined in aws-sdk-js: https://github.com/aws/aws-sdk-js/blob/02153340ce052614c5437fe46ab1f49bfd8483f7/lib/protocol/rest_json.js#L24-L41
as well as https://github.com/aws/aws-sdk-js/blob/02153340ce052614c5437fe46ab1f49bfd8483f7/lib/protocol/rest_xml.js#L10-L26

FTI-6138
---
 README.md                                  |  6 ++
 spec/02-requests/02-build_request_spec.lua | 77 +++++++++++++++++++++-
 src/resty/aws/request/build.lua            | 15 +++++
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 83db70e..6e82280 100644
--- a/README.md
+++ b/README.md
@@ -177,6 +177,12 @@ Release process:
 1. test installing the rock from LuaRocks
 
 
+### Unreleased
+
+- fix: build the request body based on payload field
+  [126](https://github.com/Kong/lua-resty-aws/pull/126)
+
+
 ### 1.5.2 (29-Jul-2024)
 
 - fix: fix sts regional endpoint injection under several cases
diff --git a/spec/02-requests/02-build_request_spec.lua b/spec/02-requests/02-build_request_spec.lua
index 6c6e55f..9e3a1b5 100644
--- a/spec/02-requests/02-build_request_spec.lua
+++ b/spec/02-requests/02-build_request_spec.lua
@@ -5,9 +5,9 @@ describe("operations protocol", function()
 
 
   local build_request
-  local operation
-  local config
-  local params
+  local operation, operation_with_payload_field
+  local config, config_with_payload
+  local params, params_with_payload
   local snapshot
   local binary_data
 
@@ -105,6 +105,37 @@ describe("operations protocol", function()
       }
     }
 
+    operation_with_payload_field = {
+      name = "PutObject",
+      http = {
+        method = "PUT",
+        requestUri = "/{Bucket}/{Key+}"
+      },
+      input = {
+        type = "structure",
+        required = {
+          "Bucket",
+          "Key"
+        },
+        members = {
+          Bucket = {
+            type = "string",
+            location = "uri",
+            locationName = "Bucket"
+          },
+          Key = {
+            type = "string",
+            location = "uri",
+            locationName = "Key"
+          },
+          Body = {
+            type = "blob",
+          },
+        },
+        payload = "Body"
+      },
+    }
+
     config = {
       apiVersion = "2011-06-15",
       --endpointPrefix = "sts",
@@ -119,6 +150,19 @@ describe("operations protocol", function()
       xmlNamespace = "https://sts.amazonaws.com/doc/2011-06-15/"
     }
 
+    config_with_payload = {
+      apiVersion = "2006-03-01",
+      signingName = "s3",
+      globalEndpoint = "s3.amazonaws.com",
+      --protocol = "query",
+      serviceAbbreviation = "AWS S3",
+      serviceFullName = "AWS Object Storage",
+      serviceId = "S3",
+      signatureVersion = "v4",
+      uid = "s3-2006-03-01",
+      xmlNamespace = "https://s3.amazonaws.com/doc/2006-03-01/"
+    }
+
     params = {
       RoleArn = "hello",
       RoleSessionName = "world",
@@ -134,6 +178,12 @@ describe("operations protocol", function()
       BinaryData = binary_data,
     }
 
+    params_with_payload = {
+      Bucket = "hello",
+      Key = "world",
+      Body = binary_data,
+    }
+
   end)
 
 
@@ -256,6 +306,27 @@ describe("operations protocol", function()
     }, request)
   end)
 
+  it("json: querystring, uri, header and body params, with payload field", function()
+
+    config_with_payload.protocol = "json"
+
+    local request = build_request(operation_with_payload_field, config_with_payload, params_with_payload)
+
+    assert.same({
+      headers = {
+        ["Content-Length"] = 4,
+        ["X-Amz-Target"] = "s3.PutObject",
+        ["Host"] = "s3.amazonaws.com",
+      },
+      method = 'PUT',
+      path = '/hello/world',
+      host = 's3.amazonaws.com',
+      port = 443,
+      body = binary_data,
+      query = {},
+    }, request)
+  end)
+
 
   it("rest-xml: querystring, uri, header and body params", function()
 
diff --git a/src/resty/aws/request/build.lua b/src/resty/aws/request/build.lua
index 22b78f6..7b5794e 100644
--- a/src/resty/aws/request/build.lua
+++ b/src/resty/aws/request/build.lua
@@ -205,6 +205,21 @@ local function build_request(operation, config, params)
     request.query[k] = v
   end
 
+  local payload_member = operation.input and operation.input.payload
+  if payload_member and body_tbl[payload_member] then
+    local member_config = operation.input.members[payload_member]
+    if member_config.type == "structure" then
+      body_tbl = type(body_tbl[payload_member]) == "table" and body_tbl[payload_member] or body_tbl
+
+    elseif body_tbl[payload_member] then
+      request.body = body_tbl[payload_member]
+      if member_config.type == "binary" and member_config.streaming == true then
+        request.headers["Content-Type"] = "binary/octet-stream"
+      end
+      body_tbl[payload_member] = nil
+    end
+  end
+
   local table_empty = not next(body_tbl)
   -- already has a raw body, or no body at all
   if request.body then