aws/client: Fix HTTP debug log EventStream payloads

aws/client/client.go

@@ -91,6 +91,6 @@ func (c *Client) AddDebugHandlers() {
 		return
 	}
 
-	c.Handlers.Send.PushFrontNamed(request.NamedHandler{Name: "awssdk.client.LogRequest", Fn: logRequest})
-	c.Handlers.Send.PushBackNamed(request.NamedHandler{Name: "awssdk.client.LogResponse", Fn: logResponse})
+	c.Handlers.Send.PushFrontNamed(LogHTTPRequestHandler)
+	c.Handlers.Send.PushBackNamed(LogHTTPResponseHandler)
 }

aws/client/logger.go

@@ -44,12 +44,22 @@ func (reader *teeReaderCloser) Close() error {
 	return reader.Source.Close()
 }
 
+// LogHTTPRequestHandler is a SDK request handler to log the HTTP request sent
+// to a service. Will include the HTTP request body if the LogLevel of the
+// request matches LogDebugWithHTTPBody.
+var LogHTTPRequestHandler = request.NamedHandler{
+	Name: "awssdk.client.LogRequest",
+	Fn:   logRequest,
+}
+
 func logRequest(r *request.Request) {
 	logBody := r.Config.LogLevel.Matches(aws.LogDebugWithHTTPBody)
 	bodySeekable := aws.IsReaderSeekable(r.Body)
-	dumpedBody, err := httputil.DumpRequestOut(r.HTTPRequest, logBody)
+
+	b, err := httputil.DumpRequestOut(r.HTTPRequest, logBody)
 	if err != nil {
-		r.Config.Logger.Log(fmt.Sprintf(logReqErrMsg, r.ClientInfo.ServiceName, r.Operation.Name, err))
+		r.Config.Logger.Log(fmt.Sprintf(logReqErrMsg,
+			r.ClientInfo.ServiceName, r.Operation.Name, err))
 		return
 	}
 
@@ -63,7 +73,28 @@ func logRequest(r *request.Request) {
 		r.ResetBody()
 	}
 
-	r.Config.Logger.Log(fmt.Sprintf(logReqMsg, r.ClientInfo.ServiceName, r.Operation.Name, string(dumpedBody)))
+	r.Config.Logger.Log(fmt.Sprintf(logReqMsg,
+		r.ClientInfo.ServiceName, r.Operation.Name, string(b)))
+}
+
+// LogHTTPRequestHeaderHandler is a SDK request handler to log the HTTP request sent
+// to a service. Will only log the HTTP request's headers. The request payload
+// will not be read.
+var LogHTTPRequestHeaderHandler = request.NamedHandler{
+	Name: "awssdk.client.LogRequestHeader",
+	Fn:   logRequestHeader,
+}
+
+func logRequestHeader(r *request.Request) {
+	b, err := httputil.DumpRequestOut(r.HTTPRequest, false)
+	if err != nil {
+		r.Config.Logger.Log(fmt.Sprintf(logReqErrMsg,
+			r.ClientInfo.ServiceName, r.Operation.Name, err))
+		return
+	}
+
+	r.Config.Logger.Log(fmt.Sprintf(logReqMsg,
+		r.ClientInfo.ServiceName, r.Operation.Name, string(b)))
 }
 
 const logRespMsg = `DEBUG: Response %s/%s Details:
@@ -76,27 +107,44 @@ const logRespErrMsg = `DEBUG ERROR: Response %s/%s:
 %s
 -----------------------------------------------------`
 
+// LogHTTPResponseHandler is a SDK request handler to log the HTTP response
+// received from a service. Will include the HTTP response body if the LogLevel
+// of the request matches LogDebugWithHTTPBody.
+var LogHTTPResponseHandler = request.NamedHandler{
+	Name: "awssdk.client.LogResponse",
+	Fn:   logResponse,
+}
+
 func logResponse(r *request.Request) {
 	lw := &logWriter{r.Config.Logger, bytes.NewBuffer(nil)}
-	r.HTTPResponse.Body = &teeReaderCloser{
-		Reader: io.TeeReader(r.HTTPResponse.Body, lw),
-		Source: r.HTTPResponse.Body,
+
+	logBody := r.Config.LogLevel.Matches(aws.LogDebugWithHTTPBody)
+	if logBody {
+		r.HTTPResponse.Body = &teeReaderCloser{
+			Reader: io.TeeReader(r.HTTPResponse.Body, lw),
+			Source: r.HTTPResponse.Body,
+		}
 	}
 
 	handlerFn := func(req *request.Request) {
-		body, err := httputil.DumpResponse(req.HTTPResponse, false)
+		b, err := httputil.DumpResponse(req.HTTPResponse, false)
 		if err != nil {
-			lw.Logger.Log(fmt.Sprintf(logRespErrMsg, req.ClientInfo.ServiceName, req.Operation.Name, err))
+			lw.Logger.Log(fmt.Sprintf(logRespErrMsg,
+				req.ClientInfo.ServiceName, req.Operation.Name, err))
 			return
 		}
 
-		b, err := ioutil.ReadAll(lw.buf)
-		if err != nil {
-			lw.Logger.Log(fmt.Sprintf(logRespErrMsg, req.ClientInfo.ServiceName, req.Operation.Name, err))
-			return
-		}
-		lw.Logger.Log(fmt.Sprintf(logRespMsg, req.ClientInfo.ServiceName, req.Operation.Name, string(body)))
-		if req.Config.LogLevel.Matches(aws.LogDebugWithHTTPBody) {
+		lw.Logger.Log(fmt.Sprintf(logRespMsg,
+			req.ClientInfo.ServiceName, req.Operation.Name, string(b)))
+
+		if logBody {
+			b, err := ioutil.ReadAll(lw.buf)
+			if err != nil {
+				lw.Logger.Log(fmt.Sprintf(logRespErrMsg,
+					req.ClientInfo.ServiceName, req.Operation.Name, err))
+				return
+			}
+
 			lw.Logger.Log(string(b))
 		}
 	}
@@ -110,3 +158,27 @@ func logResponse(r *request.Request) {
 		Name: handlerName, Fn: handlerFn,
 	})
 }
+
+// LogHTTPResponseHeaderHandler is a SDK request handler to log the HTTP
+// response received from a service. Will only log the HTTP response's headers.
+// The response payload will not be read.
+var LogHTTPResponseHeaderHandler = request.NamedHandler{
+	Name: "awssdk.client.LogResponseHeader",
+	Fn:   logResponseHeader,
+}
+
+func logResponseHeader(r *request.Request) {
+	if r.Config.Logger == nil {
+		return
+	}
+
+	b, err := httputil.DumpResponse(r.HTTPResponse, false)
+	if err != nil {
+		r.Config.Logger.Log(fmt.Sprintf(logRespErrMsg,
+			r.ClientInfo.ServiceName, r.Operation.Name, err))
+		return
+	}
+
+	r.Config.Logger.Log(fmt.Sprintf(logRespMsg,
+		r.ClientInfo.ServiceName, r.Operation.Name, string(b)))
+}

aws/client/logger_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net/http"
 	"reflect"
 	"testing"
 
@@ -127,6 +128,83 @@ func TestLogRequest(t *testing.T) {
 	}
 }
 
+func TestLogResponse(t *testing.T) {
+	cases := []struct {
+		Body       *bytes.Buffer
+		ExpectBody []byte
+		ReadBody   bool
+		LogLevel   aws.LogLevelType
+	}{
+		{
+			Body:       bytes.NewBuffer([]byte("body content")),
+			ExpectBody: []byte("body content"),
+		},
+		{
+			Body:       bytes.NewBuffer([]byte("body content")),
+			LogLevel:   aws.LogDebug,
+			ExpectBody: []byte("body content"),
+		},
+		{
+			Body:       bytes.NewBuffer([]byte("body content")),
+			LogLevel:   aws.LogDebugWithHTTPBody,
+			ReadBody:   true,
+			ExpectBody: []byte("body content"),
+		},
+	}
+
+	for i, c := range cases {
+		var logW bytes.Buffer
+		req := request.New(
+			aws.Config{
+				Credentials: credentials.AnonymousCredentials,
+				Logger:      &bufLogger{w: &logW},
+				LogLevel:    aws.LogLevel(c.LogLevel),
+			},
+			metadata.ClientInfo{
+				Endpoint: "https://mock-service.mock-region.amazonaws.com",
+			},
+			testHandlers(),
+			nil,
+			&request.Operation{
+				Name:       "APIName",
+				HTTPMethod: "POST",
+				HTTPPath:   "/",
+			},
+			struct{}{}, nil,
+		)
+		req.HTTPResponse = &http.Response{
+			StatusCode: 200,
+			Status:     "OK",
+			Header: http.Header{
+				"ABC": []string{"123"},
+			},
+			Body: ioutil.NopCloser(c.Body),
+		}
+
+		logResponse(req)
+		req.Handlers.Unmarshal.Run(req)
+
+		if c.ReadBody {
+			if e, a := len(c.ExpectBody), c.Body.Len(); e != a {
+				t.Errorf("%d, expect orginal body not to of been read", i)
+			}
+		}
+
+		if logW.Len() == 0 {
+			t.Errorf("%d, expect HTTP Response headers to be logged", i)
+		}
+
+		b, err := ioutil.ReadAll(req.HTTPResponse.Body)
+		if err != nil {
+			t.Fatalf("%d, expect to read SDK request Body", i)
+		}
+
+		if e, a := c.ExpectBody, b; !bytes.Equal(e, a) {
+			t.Errorf("%d, expect %v body, got %v", i, e, a)
+		}
+	}
+}
+
 type bufLogger struct {
 	w *bytes.Buffer
 }

private/model/api/operation.go

@@ -129,6 +129,7 @@ func (c *{{ .API.StructName }}) {{ .ExportedName }}Request(` +
 	{{ end -}}
 	{{ if ne .AuthType "" }}{{ .GetSigner }}{{ end -}}
 	{{ if .OutputRef.Shape.EventStreamsMemberName -}}
+		req.Handlers.Send.Swap(client.LogHTTPResponseHandler.Name, client.LogHTTPResponseHeaderHandler)
 		req.Handlers.Unmarshal.Swap({{ .API.ProtocolPackage }}.UnmarshalHandler.Name, rest.UnmarshalHandler)
 		req.Handlers.Unmarshal.PushBack(output.runEventStreamLoop)
 	{{ end -}}
@@ -252,6 +253,7 @@ func (o *Operation) GoCode() string {
 
 	if len(o.OutputRef.Shape.EventStreamsMemberName) != 0 {
 		// TODO need better was of updating protocol unmarshalers
+		o.API.imports["github.com/aws/aws-sdk-go/aws/client"] = true
 		o.API.imports["github.com/aws/aws-sdk-go/private/protocol/rest"] = true
 	}