accept_encoding_gzip.go raw

   1  package acceptencoding
   2  
   3  import (
   4  	"compress/gzip"
   5  	"context"
   6  	"fmt"
   7  	"io"
   8  
   9  	"github.com/aws/smithy-go"
  10  	"github.com/aws/smithy-go/middleware"
  11  	smithyhttp "github.com/aws/smithy-go/transport/http"
  12  )
  13  
  14  const acceptEncodingHeaderKey = "Accept-Encoding"
  15  const contentEncodingHeaderKey = "Content-Encoding"
  16  
  17  // AddAcceptEncodingGzipOptions provides the options for the
  18  // AddAcceptEncodingGzip middleware setup.
  19  type AddAcceptEncodingGzipOptions struct {
  20  	Enable bool
  21  }
  22  
  23  // AddAcceptEncodingGzip explicitly adds handling for accept-encoding GZIP
  24  // middleware to the operation stack. This allows checksums to be correctly
  25  // computed without disabling GZIP support.
  26  func AddAcceptEncodingGzip(stack *middleware.Stack, options AddAcceptEncodingGzipOptions) error {
  27  	if options.Enable {
  28  		if err := stack.Finalize.Add(&EnableGzip{}, middleware.Before); err != nil {
  29  			return err
  30  		}
  31  		if err := stack.Deserialize.Insert(&DecompressGzip{}, "OperationDeserializer", middleware.After); err != nil {
  32  			return err
  33  		}
  34  		return nil
  35  	}
  36  
  37  	return stack.Finalize.Add(&DisableGzip{}, middleware.Before)
  38  }
  39  
  40  // DisableGzip provides the middleware that will
  41  // disable the underlying http client automatically enabling for gzip
  42  // decompress content-encoding support.
  43  type DisableGzip struct{}
  44  
  45  // ID returns the id for the middleware.
  46  func (*DisableGzip) ID() string {
  47  	return "DisableAcceptEncodingGzip"
  48  }
  49  
  50  // HandleFinalize implements the FinalizeMiddleware interface.
  51  func (*DisableGzip) HandleFinalize(
  52  	ctx context.Context, input middleware.FinalizeInput, next middleware.FinalizeHandler,
  53  ) (
  54  	output middleware.FinalizeOutput, metadata middleware.Metadata, err error,
  55  ) {
  56  	req, ok := input.Request.(*smithyhttp.Request)
  57  	if !ok {
  58  		return output, metadata, &smithy.SerializationError{
  59  			Err: fmt.Errorf("unknown request type %T", input.Request),
  60  		}
  61  	}
  62  
  63  	// Explicitly enable gzip support, this will prevent the http client from
  64  	// auto extracting the zipped content.
  65  	req.Header.Set(acceptEncodingHeaderKey, "identity")
  66  
  67  	return next.HandleFinalize(ctx, input)
  68  }
  69  
  70  // EnableGzip provides a middleware to enable support for
  71  // gzip responses, with manual decompression. This prevents the underlying HTTP
  72  // client from performing the gzip decompression automatically.
  73  type EnableGzip struct{}
  74  
  75  // ID returns the id for the middleware.
  76  func (*EnableGzip) ID() string {
  77  	return "AcceptEncodingGzip"
  78  }
  79  
  80  // HandleFinalize implements the FinalizeMiddleware interface.
  81  func (*EnableGzip) HandleFinalize(
  82  	ctx context.Context, input middleware.FinalizeInput, next middleware.FinalizeHandler,
  83  ) (
  84  	output middleware.FinalizeOutput, metadata middleware.Metadata, err error,
  85  ) {
  86  	req, ok := input.Request.(*smithyhttp.Request)
  87  	if !ok {
  88  		return output, metadata, &smithy.SerializationError{
  89  			Err: fmt.Errorf("unknown request type %T", input.Request),
  90  		}
  91  	}
  92  
  93  	// Explicitly enable gzip support, this will prevent the http client from
  94  	// auto extracting the zipped content.
  95  	req.Header.Set(acceptEncodingHeaderKey, "gzip")
  96  
  97  	return next.HandleFinalize(ctx, input)
  98  }
  99  
 100  // DecompressGzip provides the middleware for decompressing a gzip
 101  // response from the service.
 102  type DecompressGzip struct{}
 103  
 104  // ID returns the id for the middleware.
 105  func (*DecompressGzip) ID() string {
 106  	return "DecompressGzip"
 107  }
 108  
 109  // HandleDeserialize implements the DeserializeMiddlware interface.
 110  func (*DecompressGzip) HandleDeserialize(
 111  	ctx context.Context, input middleware.DeserializeInput, next middleware.DeserializeHandler,
 112  ) (
 113  	output middleware.DeserializeOutput, metadata middleware.Metadata, err error,
 114  ) {
 115  	output, metadata, err = next.HandleDeserialize(ctx, input)
 116  	if err != nil {
 117  		return output, metadata, err
 118  	}
 119  
 120  	resp, ok := output.RawResponse.(*smithyhttp.Response)
 121  	if !ok {
 122  		return output, metadata, &smithy.DeserializationError{
 123  			Err: fmt.Errorf("unknown response type %T", output.RawResponse),
 124  		}
 125  	}
 126  	if v := resp.Header.Get(contentEncodingHeaderKey); v != "gzip" {
 127  		return output, metadata, err
 128  	}
 129  
 130  	// Clear content length since it will no longer be valid once the response
 131  	// body is decompressed.
 132  	resp.Header.Del("Content-Length")
 133  	resp.ContentLength = -1
 134  
 135  	resp.Body = wrapGzipReader(resp.Body)
 136  
 137  	return output, metadata, err
 138  }
 139  
 140  type gzipReader struct {
 141  	reader io.ReadCloser
 142  	gzip   *gzip.Reader
 143  }
 144  
 145  func wrapGzipReader(reader io.ReadCloser) *gzipReader {
 146  	return &gzipReader{
 147  		reader: reader,
 148  	}
 149  }
 150  
 151  // Read wraps the gzip reader around the underlying io.Reader to extract the
 152  // response bytes on the fly.
 153  func (g *gzipReader) Read(b []byte) (n int, err error) {
 154  	if g.gzip == nil {
 155  		g.gzip, err = gzip.NewReader(g.reader)
 156  		if err != nil {
 157  			g.gzip = nil // ensure uninitialized gzip value isn't used in close.
 158  			return 0, fmt.Errorf("failed to decompress gzip response, %w", err)
 159  		}
 160  	}
 161  
 162  	return g.gzip.Read(b)
 163  }
 164  
 165  func (g *gzipReader) Close() error {
 166  	if g.gzip == nil {
 167  		return nil
 168  	}
 169  
 170  	if err := g.gzip.Close(); err != nil {
 171  		g.reader.Close()
 172  		return fmt.Errorf("failed to decompress gzip response, %w", err)
 173  	}
 174  
 175  	return g.reader.Close()
 176  }
 177