From 43d7d998b5818a71879ab45ee5e0503ef9db47b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Carlos=20Ch=C3=A1vez?= <jcchavezs@gmail.com>
Date: Wed, 22 May 2019 09:22:01 +0200
Subject: [PATCH] feat(#126): adds body response reader for better errors.

---
 middleware/http/transport.go      | 48 +++++++++++++++++++++++++++----
 middleware/http/transport_test.go | 36 +++++++++++++++++++++++
 2 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/middleware/http/transport.go b/middleware/http/transport.go
index 55d9cb2..0d2ec2d 100644
--- a/middleware/http/transport.go
+++ b/middleware/http/transport.go
@@ -15,8 +15,13 @@
 package http
 
 import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"log"
 	"net/http"
 	"net/http/httptrace"
+	"os"
 	"strconv"
 
 	zipkin "github.com/openzipkin/zipkin-go"
@@ -39,12 +44,19 @@ func defaultErrHandler(sp zipkin.Span, err error, statusCode int) {
 	zipkin.TagError.Set(sp, statusCodeVal)
 }
 
+// ErrResponseReader allows instrumentations to read the error body
+// and decide to obtain information to it and add it to the span i.e.
+// tag the span with a more meaningful error code or with error details.
+type ErrResponseReader func(sp zipkin.Span, body io.Reader)
+
 type transport struct {
-	tracer      *zipkin.Tracer
-	rt          http.RoundTripper
-	httpTrace   bool
-	defaultTags map[string]string
-	errHandler  ErrHandler
+	tracer            *zipkin.Tracer
+	rt                http.RoundTripper
+	httpTrace         bool
+	defaultTags       map[string]string
+	errHandler        ErrHandler
+	errResponseReader *ErrResponseReader
+	logger            *log.Logger
 }
 
 // TransportOption allows one to configure optional transport configuration.
@@ -80,6 +92,20 @@ func TransportErrHandler(h ErrHandler) TransportOption {
 	}
 }
 
+// TransportErrResponseReader allows to pass a custom ErrResponseReader
+func TransportErrResponseReader(r ErrResponseReader) TransportOption {
+	return func(t *transport) {
+		t.errResponseReader = &r
+	}
+}
+
+// TransportLogger allows to plug a logger into the transport
+func TransportLogger(l *log.Logger) TransportOption {
+	return func(t *transport) {
+		t.logger = l
+	}
+}
+
 // NewTransport returns a new Zipkin instrumented http RoundTripper which can be
 // used with a standard library http Client.
 func NewTransport(tracer *zipkin.Tracer, options ...TransportOption) (http.RoundTripper, error) {
@@ -92,6 +118,7 @@ func NewTransport(tracer *zipkin.Tracer, options ...TransportOption) (http.Round
 		rt:         http.DefaultTransport,
 		httpTrace:  false,
 		errHandler: defaultErrHandler,
+		logger:     log.New(os.Stderr, "", log.LstdFlags),
 	}
 
 	for _, option := range options {
@@ -157,6 +184,17 @@ func (t *transport) RoundTrip(req *http.Request) (res *http.Response, err error)
 		zipkin.TagHTTPStatusCode.Set(sp, statusCode)
 		if res.StatusCode > 399 {
 			t.errHandler(sp, nil, res.StatusCode)
+
+			if t.errResponseReader != nil {
+				sBody, err := ioutil.ReadAll(res.Body)
+				if err == nil {
+					res.Body.Close()
+					(*t.errResponseReader)(sp, ioutil.NopCloser(bytes.NewBuffer(sBody)))
+					res.Body = ioutil.NopCloser(bytes.NewBuffer(sBody))
+				} else {
+					t.logger.Printf("failed to read the response body in the ErrResponseReader: %v", err)
+				}
+			}
 		}
 	}
 	sp.Finish()
diff --git a/middleware/http/transport_test.go b/middleware/http/transport_test.go
index 76212a4..4bf3b31 100644
--- a/middleware/http/transport_test.go
+++ b/middleware/http/transport_test.go
@@ -2,6 +2,8 @@ package http
 
 import (
 	"errors"
+	"io"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -89,3 +91,37 @@ func TestRoundTripErrHandlingForStatusCode(t *testing.T) {
 		srv.Close()
 	}
 }
+
+func TestRoundTripErrResponseReadingSuccess(t *testing.T) {
+	expectedBody := []byte("message")
+	srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
+		rw.WriteHeader(500)
+		rw.Write(expectedBody)
+	}))
+	defer srv.Close()
+
+	tracer, err := zipkin.NewTracer(nil)
+	if err != nil {
+		t.Fatalf("unexpected error when creating tracer: %v", err)
+	}
+	req, _ := http.NewRequest("GET", srv.URL, nil)
+	transport, _ := NewTransport(
+		tracer,
+		TransportErrResponseReader(func(_ zipkin.Span, br io.Reader) {
+			body, _ := ioutil.ReadAll(br)
+			if want, have := expectedBody, body; string(want) != string(have) {
+				t.Errorf("unexpected body, want %q, have %q", want, have)
+			}
+		}),
+	)
+
+	res, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	actualBody, _ := ioutil.ReadAll(res.Body)
+	if want, have := expectedBody, actualBody; string(expectedBody) != string(actualBody) {
+		t.Errorf("unexpected body: want %s, have %s", want, have)
+	}
+}