From 1aaad14b05939697fa0c678d4c90cdfa92a6005c Mon Sep 17 00:00:00 2001
From: Bryan Moffatt <moffattb@amazon.com>
Date: Wed, 24 Jan 2024 09:39:38 -0800
Subject: [PATCH] fix xrayException shape - the json tags differ from
 InvokeResponse_Error

---
 lambda/invoke_loop.go      | 29 +++++++++++++++++++----------
 lambda/invoke_loop_test.go | 16 +++++++++-------
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/lambda/invoke_loop.go b/lambda/invoke_loop.go
index a121773d..e558ea45 100644
--- a/lambda/invoke_loop.go
+++ b/lambda/invoke_loop.go
@@ -102,11 +102,10 @@ func handleInvoke(invoke *invoke, handler *handlerOptions) error {
 }
 
 func reportFailure(invoke *invoke, invokeErr *messages.InvokeResponse_Error) error {
-	errorForXRay := makeErrorForXRay(invokeErr)
-	errorPayload := errorForXRay.Exceptions[0]
+	errorPayload := safeMarshal(invokeErr)
 	log.Printf("%s", errorPayload)
 
-	causeForXRay, err := json.Marshal(errorForXRay)
+	causeForXRay, err := json.Marshal(makeXRayError(invokeErr))
 	if err != nil {
 		return fmt.Errorf("unexpected error occured when serializing the function error cause for X-Ray: %v", err)
 	}
@@ -174,13 +173,19 @@ func safeMarshal(v interface{}) []byte {
 	return payload
 }
 
-type errorForXRay struct {
-	WorkingDirectory string            `json:"working_directory"`
-	Exceptions       []json.RawMessage `json:"exceptions"` // returned as bytes to avoid double-serializing
-	Paths            []string          `json:"paths"`
+type xrayException struct {
+	Type    string                                      `json:"type"`
+	Message string                                      `json:"message"`
+	Stack   []*messages.InvokeResponse_Error_StackFrame `json:"stack"`
 }
 
-func makeErrorForXRay(invokeResponseError *messages.InvokeResponse_Error) *errorForXRay {
+type xrayError struct {
+	WorkingDirectory string          `json:"working_directory"`
+	Exceptions       []xrayException `json:"exceptions"`
+	Paths            []string        `json:"paths"`
+}
+
+func makeXRayError(invokeResponseError *messages.InvokeResponse_Error) *xrayError {
 	pathSet := make(map[string]struct{}, len(invokeResponseError.StackTrace))
 	for _, frame := range invokeResponseError.StackTrace {
 		pathSet[frame.Path] = struct{}{}
@@ -190,9 +195,13 @@ func makeErrorForXRay(invokeResponseError *messages.InvokeResponse_Error) *error
 		paths = append(paths, path)
 	}
 	cwd, _ := os.Getwd()
-	return &errorForXRay{
+	return &xrayError{
 		WorkingDirectory: cwd,
 		Paths:            paths,
-		Exceptions:       []json.RawMessage{safeMarshal(invokeResponseError)},
+		Exceptions: []xrayException{{
+			Type:    invokeResponseError.Type,
+			Message: invokeResponseError.Message,
+			Stack:   invokeResponseError.StackTrace,
+		}},
 	}
 }
diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go
index ab110c50..a3a2fcd0 100644
--- a/lambda/invoke_loop_test.go
+++ b/lambda/invoke_loop_test.go
@@ -93,10 +93,11 @@ func TestCustomErrorMarshaling(t *testing.T) {
 func TestXRayCausePlumbing(t *testing.T) {
 	errors := []error{
 		messages.InvokeResponse_Error{
-			Type:    "yolo",
-			Message: "hello",
+			Type:    "yoloError",
+			Message: "hello yolo",
 			StackTrace: []*messages.InvokeResponse_Error_StackFrame{
 				{Label: "yolo", Path: "yolo", Line: 2},
+				{Label: "hi", Path: "hello/hello", Line: 12},
 			},
 		},
 	}
@@ -104,12 +105,13 @@ func TestXRayCausePlumbing(t *testing.T) {
 	expected := []string{
 		`{
 		    "working_directory":"` + wd + `", 
-		    "paths": ["yolo"], 
+		    "paths": ["yolo", "hello/hello"], 
 		    "exceptions": [{ 
-			"errorType": "yolo", 
-			"errorMessage": "hello", 
-			"stackTrace": [
-			    {"label": "yolo", "path": "yolo", "line": 2}
+			"type": "yoloError", 
+			"message": "hello yolo", 
+			"stack": [
+			    {"label": "yolo", "path": "yolo", "line": 2},
+			    {"label": "hi", "path": "hello/hello", "line": 12}
 			]
 		    }]
 		}`,