src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

middleware_close_response_body.go (2597B)

      1 package http
      2 
      3 import (
      4 	"context"
      5 	"github.com/aws/smithy-go/logging"
      6 	"github.com/aws/smithy-go/middleware"
      7 	"io"
      8 	"io/ioutil"
      9 )
     10 
     11 // AddErrorCloseResponseBodyMiddleware adds the middleware to automatically
     12 // close the response body of an operation request if the request response
     13 // failed.
     14 func AddErrorCloseResponseBodyMiddleware(stack *middleware.Stack) error {
     15 	return stack.Deserialize.Insert(&errorCloseResponseBodyMiddleware{}, "OperationDeserializer", middleware.Before)
     16 }
     17 
     18 type errorCloseResponseBodyMiddleware struct{}
     19 
     20 func (*errorCloseResponseBodyMiddleware) ID() string {
     21 	return "ErrorCloseResponseBody"
     22 }
     23 
     24 func (m *errorCloseResponseBodyMiddleware) HandleDeserialize(
     25 	ctx context.Context, input middleware.DeserializeInput, next middleware.DeserializeHandler,
     26 ) (
     27 	output middleware.DeserializeOutput, metadata middleware.Metadata, err error,
     28 ) {
     29 	out, metadata, err := next.HandleDeserialize(ctx, input)
     30 	if err != nil {
     31 		if resp, ok := out.RawResponse.(*Response); ok && resp != nil && resp.Body != nil {
     32 			// Consume the full body to prevent TCP connection resets on some platforms
     33 			_, _ = io.Copy(ioutil.Discard, resp.Body)
     34 			// Do not validate that the response closes successfully.
     35 			resp.Body.Close()
     36 		}
     37 	}
     38 
     39 	return out, metadata, err
     40 }
     41 
     42 // AddCloseResponseBodyMiddleware adds the middleware to automatically close
     43 // the response body of an operation request, after the response had been
     44 // deserialized.
     45 func AddCloseResponseBodyMiddleware(stack *middleware.Stack) error {
     46 	return stack.Deserialize.Insert(&closeResponseBody{}, "OperationDeserializer", middleware.Before)
     47 }
     48 
     49 type closeResponseBody struct{}
     50 
     51 func (*closeResponseBody) ID() string {
     52 	return "CloseResponseBody"
     53 }
     54 
     55 func (m *closeResponseBody) HandleDeserialize(
     56 	ctx context.Context, input middleware.DeserializeInput, next middleware.DeserializeHandler,
     57 ) (
     58 	output middleware.DeserializeOutput, metadata middleware.Metadata, err error,
     59 ) {
     60 	out, metadata, err := next.HandleDeserialize(ctx, input)
     61 	if err != nil {
     62 		return out, metadata, err
     63 	}
     64 
     65 	if resp, ok := out.RawResponse.(*Response); ok {
     66 		// Consume the full body to prevent TCP connection resets on some platforms
     67 		_, copyErr := io.Copy(ioutil.Discard, resp.Body)
     68 		if copyErr != nil {
     69 			middleware.GetLogger(ctx).Logf(logging.Warn, "failed to discard remaining HTTP response body, this may affect connection reuse")
     70 		}
     71 
     72 		closeErr := resp.Body.Close()
     73 		if closeErr != nil {
     74 			middleware.GetLogger(ctx).Logf(logging.Warn, "failed to close HTTP response body, this may affect connection reuse")
     75 		}
     76 	}
     77 
     78 	return out, metadata, err
     79 }