From 8a710e49232e4bc29f37bef6796ac60402a450c8 Mon Sep 17 00:00:00 2001 From: AndreasW Date: Sun, 10 Mar 2019 15:32:11 +0100 Subject: [PATCH 1/2] Flag for disable response HTML encoding while json marshaling lambda response --- lambda/handler.go | 31 ++++++++++++++++++++++++ lambda/handler_test.go | 54 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/lambda/handler.go b/lambda/handler.go index 96eeac2e..81feb986 100644 --- a/lambda/handler.go +++ b/lambda/handler.go @@ -3,10 +3,12 @@ package lambda import ( + "bytes" "context" "encoding/json" "fmt" "reflect" + "sync" "github.com/aws/aws-lambda-go/lambda/handlertrace" ) @@ -18,6 +20,9 @@ type Handler interface { // lambdaHandler is the generic function type type lambdaHandler func(context.Context, []byte) (interface{}, error) +var responseDisableEscapeHTML = false +var muresponseDisableEscapeHTML sync.Mutex + // Invoke calls the handler, and serializes the response. // If the underlying handler returned an error, or an error occurs during serialization, error is returned. func (handler lambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) { @@ -26,6 +31,10 @@ func (handler lambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte return nil, err } + if responseDisableEscapeHTML { + return MarshalWithDisabledEscapeHTML(response) + } + responseBytes, err := json.Marshal(response) if err != nil { return nil, err @@ -34,6 +43,19 @@ func (handler lambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte return responseBytes, nil } +//MarshalWithDisabledEscapeHTML uses json.NewEncoder instead of json.Marshal to disable escaping HTML since it may corrupt response. +func MarshalWithDisabledEscapeHTML(response interface{}) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + err := enc.Encode(response) + if err != nil { + return nil, err + } + + return buf.Bytes(), nil +} + func errorHandler(e error) lambdaHandler { return func(ctx context.Context, event []byte) (interface{}, error) { return nil, e @@ -72,6 +94,15 @@ func validateReturns(handler reflect.Type) error { return nil } +// SetResponseNonEscapeHTML will NOT coerced JSON strings to valid UTF-8, +// The angle brackets "<" and ">" will NOT be escaped to "\u003c" and "\u003e" +// Ampersand "&" is also NOT be escaped to "\u0026". +func SetResponseNonEscapeHTML(flag bool) { + muresponseDisableEscapeHTML.Lock() + responseDisableEscapeHTML = flag + muresponseDisableEscapeHTML.Unlock() +} + // NewHandler creates a base lambda handler from the given handler function. The // returned Handler performs JSON deserialization and deserialization, and // delegates to the input handler function. The handler function parameter must diff --git a/lambda/handler_test.go b/lambda/handler_test.go index 962ac38c..ea6deee0 100644 --- a/lambda/handler_test.go +++ b/lambda/handler_test.go @@ -201,6 +201,60 @@ func TestInvokes(t *testing.T) { } } +func TestInvokeResponseStructs(t *testing.T) { + type testStructPerson = struct { + Name string + Age int64 + Active bool + } + + type testStructURL = struct { + URL string + } + testCases := []struct { + name string + disableEscapeHTML bool + expected expected + handler interface{} + }{ + { + name: "basic input struct serialization ", + expected: expected{`{"Name":"Aws","Age":64,"Active":true}`, nil}, + handler: func() (testStructPerson, error) { + return testStructPerson{"Aws", 64, true}, nil + }, + }, + { + name: "basic input struct serialization URL", + disableEscapeHTML: true, + expected: expected{`{"URL":"https://github.com/aws/aws-lambda-go?hello=world&go=lang"}` + "\n", nil}, + handler: func() (testStructURL, error) { + return testStructURL{URL: `https://github.com/aws/aws-lambda-go?hello=world&go=lang`}, nil + }, + }, + { + name: "basic input struct serialization URL", + expected: expected{`{"URL":"https://github.com/aws/aws-lambda-go?hello=world\u0026go=lang"}`, nil}, + handler: func() (testStructURL, error) { + return testStructURL{URL: `https://github.com/aws/aws-lambda-go?hello=world&go=lang`}, nil + }, + }, + } + for i, testCase := range testCases { + t.Run(fmt.Sprintf("testCase[%d] %s", i, testCase.name), func(t *testing.T) { + lambdaHandler := NewHandler(testCase.handler) + SetResponseNonEscapeHTML(testCase.disableEscapeHTML) + response, err := lambdaHandler.Invoke(context.TODO(), []byte(`"Lambda"`)) + if testCase.expected.err != nil { + assert.Equal(t, testCase.expected.err, err) + } else { + assert.NoError(t, err) + assert.Equal(t, testCase.expected.val, string(response)) + } + }) + } +} + func TestInvalidJsonInput(t *testing.T) { lambdaHandler := NewHandler(func(s string) error { return nil }) _, err := lambdaHandler.Invoke(context.TODO(), []byte(`{"invalid json`)) From 1f4ec0538c361f347fc467f138e46d43ced19d2c Mon Sep 17 00:00:00 2001 From: AndreasW Date: Sun, 10 Mar 2019 15:40:47 +0100 Subject: [PATCH 2/2] change marshalWithDisabledEscapeHTML to private --- lambda/handler.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lambda/handler.go b/lambda/handler.go index 81feb986..1fa511c5 100644 --- a/lambda/handler.go +++ b/lambda/handler.go @@ -32,7 +32,7 @@ func (handler lambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte } if responseDisableEscapeHTML { - return MarshalWithDisabledEscapeHTML(response) + return marshalWithDisabledEscapeHTML(response) } responseBytes, err := json.Marshal(response) @@ -43,8 +43,8 @@ func (handler lambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte return responseBytes, nil } -//MarshalWithDisabledEscapeHTML uses json.NewEncoder instead of json.Marshal to disable escaping HTML since it may corrupt response. -func MarshalWithDisabledEscapeHTML(response interface{}) ([]byte, error) { +//marshalWithDisabledEscapeHTML uses json.NewEncoder instead of json.Marshal to disable escaping HTML since it may corrupt response. +func marshalWithDisabledEscapeHTML(response interface{}) ([]byte, error) { var buf bytes.Buffer enc := json.NewEncoder(&buf) enc.SetEscapeHTML(false)