From 4f199c6f7e88832e9976f8948486b9d8d347f961 Mon Sep 17 00:00:00 2001
From: Ichinose Shogo <shogo82148@gmail.com>
Date: Sat, 19 Sep 2020 19:57:10 +0900
Subject: [PATCH 1/2] safeMarshal should escape the error message

---
 lambda/invoke_loop.go | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/lambda/invoke_loop.go b/lambda/invoke_loop.go
index 25fc6fa4..83cd05df 100644
--- a/lambda/invoke_loop.go
+++ b/lambda/invoke_loop.go
@@ -13,9 +13,8 @@ import (
 )
 
 const (
-	serializationErrorFormat = `{"errorType": "Runtime.SerializationError", "errorMessage": "%s"}`
-	msPerS                   = int64(time.Second / time.Millisecond)
-	nsPerMS                  = int64(time.Millisecond / time.Nanosecond)
+	msPerS  = int64(time.Second / time.Millisecond)
+	nsPerMS = int64(time.Millisecond / time.Nanosecond)
 )
 
 // startRuntimeAPILoop will return an error if handling a particular invoke resulted in a non-recoverable error
@@ -103,7 +102,15 @@ func convertInvokeRequest(invoke *invoke) (*messages.InvokeRequest, error) {
 func safeMarshal(v interface{}) []byte {
 	payload, err := json.Marshal(v)
 	if err != nil {
-		return []byte(fmt.Sprintf(serializationErrorFormat, err.Error()))
+		v := &messages.InvokeResponse_Error{
+			Type:    "Runtime.SerializationError",
+			Message: err.Error(),
+		}
+		payload, err := json.Marshal(v)
+		if err != nil {
+			panic(err) // never reach
+		}
+		return payload
 	}
 	return payload
 }

From ea8c3c5284509b81c3dc50b6d219fed2760e493f Mon Sep 17 00:00:00 2001
From: Ichinose Shogo <shogo82148@gmail.com>
Date: Fri, 25 Sep 2020 08:08:11 +0900
Subject: [PATCH 2/2] add a test for safeMarshal

---
 lambda/invoke_loop_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go
index 9af51874..8e76ac2f 100644
--- a/lambda/invoke_loop_test.go
+++ b/lambda/invoke_loop_test.go
@@ -106,6 +106,18 @@ func TestReadPayload(t *testing.T) {
 
 }
 
+type invalidPayload struct{}
+
+func (invalidPayload) MarshalJSON() ([]byte, error) {
+	return nil, errors.New(`some error that contains '"'`)
+}
+
+func TestSafeMarshal_SerializationError(t *testing.T) {
+	payload := safeMarshal(invalidPayload{})
+	want := `{"errorMessage":"json: error calling MarshalJSON for type lambda.invalidPayload: some error that contains '\"'","errorType":"Runtime.SerializationError"}`
+	assert.Equal(t, want, string(payload))
+}
+
 type requestRecord struct {
 	nGets     int
 	nPosts    int