From 642529985b3c3ce0d3b11ed4692adf47907c5bfc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matej=20Balanti=C4=8D?= <matej@balantic.si>
Date: Wed, 23 Oct 2019 21:49:43 +0200
Subject: [PATCH] Handle multiple returns, bubble up lambda panic

Firstly this fixes a bug where if lambda handler has multiple return types (such as `(string, error)`) and something inside lambda handler panicked, an unexpected runtime error would occur with message "*reflect: wrong return count from function created by MakeFunc*". That's because the panic was being swallowed by the wrapper, and return values weren't set because the child handler never completed (or returned). Thus, wrapper returned with the wrong number of return values (zero), triggering this runtime issue.

Secondly I am changing the behaviour of lambda wrapper so that the panic insider the handler is no longer swallowed, but rather bubbled back up. It's [perfectly valid strategy][1] to panic inside the lambda handler, and if that occurs, error should be bubbled up to AWS SDK (after it's reported to Rollbar) so that the response to the lambda request can be correctly indicated as failed, and it can get reported to CloudWatch.

Thirdly, this fixes a potentially problematic bug where `c.LogPanic` would be called at the end of each handler, regardles of whether or not there was a panic. It's currently not an issue as `LogPanic` is a NO-OP when `err` is nil, but there's no guarantee for that in the future.

[1]: https://docs.aws.amazon.com/lambda/latest/dg/go-programming-model-errors.html#go-errors-panic
---
 client.go      |  5 +++-
 client_test.go | 64 +++++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/client.go b/client.go
index 84b4f45..d083a80 100644
--- a/client.go
+++ b/client.go
@@ -554,7 +554,10 @@ func (c *Client) LambdaWrapper(handlerFunc interface{}) interface{} {
 	handler := func(args []reflect.Value) []reflect.Value {
 		defer func() {
 			err := recover()
-			c.LogPanic(err, true)
+			if err != nil {
+				c.LogPanic(err, true)
+				panic(err)
+			}
 		}()
 
 		ret := handlerValue.Call(args)
diff --git a/client_test.go b/client_test.go
index 7eaf3eb..7e00f62 100644
--- a/client_test.go
+++ b/client_test.go
@@ -283,6 +283,30 @@ type TestMessage struct {
 func TestLambdaWrapperWithError(t *testing.T) {
 	client := testClient()
 	err := errors.New("bork")
+
+	defer func() {
+		recoveredError := recover()
+
+		if recoveredError != err {
+			if recoveredError == nil {
+				t.Error("Expected wrapper to bubble up the custom panic error")
+			} else {
+				t.Errorf("Unexpected panic %s", recoveredError)
+			}
+		}
+
+		if transport, ok := client.Transport.(*TestTransport); ok {
+			if transport.Body == nil {
+				t.Error("Expected Body to be present")
+			}
+			if !transport.WaitCalled {
+				t.Error("Expected wait to be called")
+			}
+		} else {
+			t.Fail()
+		}
+	}()
+
 	//ctx := context.TODO()
 	handler := client.LambdaWrapper(func() {
 		panic(err)
@@ -291,17 +315,41 @@ func TestLambdaWrapperWithError(t *testing.T) {
 	var args []reflect.Value
 	fn.Call(args)
 	//testCallLambdaHandler(handler)
+}
 
-	if transport, ok := client.Transport.(*TestTransport); ok {
-		if transport.Body == nil {
-			t.Error("Expected Body to be present")
+func TestLambdaWrapperWithErrorAndMultipleReturnValues(t *testing.T) {
+	client := testClient()
+	err := errors.New("bork")
+
+	defer func() {
+		recoveredError := recover()
+
+		if recoveredError != err {
+			if recoveredError == nil {
+				t.Error("Expected wrapper to bubble up the custom panic error")
+			} else {
+				t.Errorf("Unexpected panic %s", recoveredError)
+			}
 		}
-		if !transport.WaitCalled {
-			t.Error("Expected wait to be called")
+
+		if transport, ok := client.Transport.(*TestTransport); ok {
+			if transport.Body == nil {
+				t.Error("Expected Body to be present")
+			}
+			if !transport.WaitCalled {
+				t.Error("Expected wait to be called")
+			}
+		} else {
+			t.Fail()
 		}
-	} else {
-		t.Fail()
-	}
+	}()
+
+	handler := client.LambdaWrapper(func() (string, error) {
+		panic(err)
+	})
+	fn := reflect.ValueOf(handler)
+	var args []reflect.Value
+	fn.Call(args)
 }
 
 func TestLambdaWrapperWithContext(t *testing.T) {