func main() {
	lambda.StartHandler(middleware(handler))
}

type handlerFunc func(ctx context.Context, payload []byte) ([]byte, error)

func (h handlerFunc) Invoke(ctx context.Context, input []byte) ([]byte, error) {
	log.Printf("[DEBUG] input = %s", string(input))
	output, err := h(ctx, input)
	log.Printf("[DEBUG] output = %s", string(output))
	return output, err
}

func middleware(originalHandler interface{}) lambda.Handler {
	handler := lambda.NewHandler(originalHandler)
	return handlerFunc(func(ctx context.Context, payload []byte) ([]byte, error) {
		ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{
			RequestEvent: func(c context.Context, e interface{}) {
				log.Printf("[DEBUG] input = %+v", e)
			},
			ResponseEvent: func(c context.Context, e interface{}) {
				log.Printf("[DEBUG] output = %+v", e)
			},
		})
		return handler.Invoke(ctx, payload)
	})
}

func handler(input Input) (*Output, error) {
  // ...
}