From f45da41b330b5a2eecf70f2ffd302267c5999a34 Mon Sep 17 00:00:00 2001
From: Bert Blommers <bblommers@users.noreply.github.com>
Date: Fri, 14 Jan 2022 22:23:09 -0100
Subject: [PATCH] AWSLambda - throw error on large responses (>6MB) (#4763)

---
 moto/awslambda/responses.py                | 11 +++++++
 tests/test_awslambda/test_lambda_invoke.py | 34 ++++++++++++++++++++++
 tests/test_awslambda/utilities.py          |  9 ++++++
 3 files changed, 54 insertions(+)

diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py
index a0be74764..eabc7e4c9 100644
--- a/moto/awslambda/responses.py
+++ b/moto/awslambda/responses.py
@@ -1,4 +1,5 @@
 import json
+import sys
 
 try:
     from urllib import unquote
@@ -207,6 +208,16 @@ class LambdaResponse(BaseResponse):
             function_name, qualifier, self.body, self.headers, response_headers
         )
         if payload:
+            if request.headers.get("X-Amz-Invocation-Type") != "Event":
+                if sys.getsizeof(payload) > 6000000:
+                    response_headers["Content-Length"] = "142"
+                    response_headers["x-amz-function-error"] = "Unhandled"
+                    error_dict = {
+                        "errorMessage": "Response payload size exceeded maximum allowed payload size (6291556 bytes).",
+                        "errorType": "Function.ResponseSizeTooLarge",
+                    }
+                    payload = json.dumps(error_dict).encode("utf-8")
+
             response_headers["content-type"] = "application/json"
             if request.headers.get("X-Amz-Invocation-Type") == "Event":
                 status_code = 202
diff --git a/tests/test_awslambda/test_lambda_invoke.py b/tests/test_awslambda/test_lambda_invoke.py
index fc7189d4a..8db0e9ef2 100644
--- a/tests/test_awslambda/test_lambda_invoke.py
+++ b/tests/test_awslambda/test_lambda_invoke.py
@@ -19,6 +19,7 @@ from .utilities import (
     get_test_zip_file1,
     get_zip_with_multiple_files,
     get_test_zip_file2,
+    get_test_zip_largeresponse,
 )
 
 _lambda_region = "us-west-2"
@@ -294,3 +295,36 @@ def test_invoke_async_function(key):
     )
 
     success_result["Status"].should.equal(202)
+
+
+@pytest.mark.network
+@mock_lambda
+def test_invoke_function_large_response():
+    # AWS Lambda should only return bodies smaller than 6 MB
+    conn = boto3.client("lambda", _lambda_region)
+    fxn = conn.create_function(
+        FunctionName=str(uuid4())[0:6],
+        Runtime="python3.7",
+        Role=get_role_name(),
+        Handler="lambda_function.lambda_handler",
+        Code={"ZipFile": get_test_zip_largeresponse()},
+        Description="test lambda function",
+        Timeout=3,
+        MemorySize=128,
+        Publish=True,
+    )
+
+    resp = conn.invoke(FunctionName=fxn["FunctionArn"])
+    resp.should.have.key("FunctionError").equals("Unhandled")
+    payload = resp["Payload"].read().decode("utf-8")
+    payload = json.loads(payload)
+    payload.should.equal(
+        {
+            "errorMessage": "Response payload size exceeded maximum allowed payload size (6291556 bytes).",
+            "errorType": "Function.ResponseSizeTooLarge",
+        }
+    )
+
+    # Absolutely fine when invoking async
+    resp = conn.invoke(FunctionName=fxn["FunctionArn"], InvocationType="Event")
+    resp.shouldnt.have.key("FunctionError")
diff --git a/tests/test_awslambda/utilities.py b/tests/test_awslambda/utilities.py
index 4b941397c..827c63082 100644
--- a/tests/test_awslambda/utilities.py
+++ b/tests/test_awslambda/utilities.py
@@ -66,6 +66,15 @@ def lambda_handler(event, context):
     return _process_lambda(pfunc)
 
 
+def get_test_zip_largeresponse():
+    pfunc = """
+def lambda_handler(event, context):
+    x = ["xxx" for x in range(10 ** 6)]
+    return {"statusCode": 200, "body": x}
+"""
+    return _process_lambda(pfunc)
+
+
 def get_zip_with_multiple_files():
     pfunc = """
 from utilities import util_function