safeMarshal should escape the error message (#327)

* safeMarshal should escape the error message

* add a test for safeMarshal

Co-authored-by: Bryan Moffatt <bmoffatt@users.noreply.github.com>

Author: shogo82148

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
 }

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