src

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

request_compression.go (3331B)

      1 // Package requestcompression implements runtime support for smithy-modeled
      2 // request compression.
      3 //
      4 // This package is designated as private and is intended for use only by the
      5 // smithy client runtime. The exported API therein is not considered stable and
      6 // is subject to breaking changes without notice.
      7 package requestcompression
      8 
      9 import (
     10 	"bytes"
     11 	"context"
     12 	"fmt"
     13 	"io"
     14 
     15 	"github.com/aws/smithy-go/middleware"
     16 	"github.com/aws/smithy-go/transport/http"
     17 )
     18 
     19 const MaxRequestMinCompressSizeBytes = 10485760
     20 
     21 // Enumeration values for supported compress Algorithms.
     22 const (
     23 	GZIP = "gzip"
     24 )
     25 
     26 type compressFunc func(io.Reader) ([]byte, error)
     27 
     28 var allowedAlgorithms = map[string]compressFunc{
     29 	GZIP: gzipCompress,
     30 }
     31 
     32 // AddRequestCompression add requestCompression middleware to op stack
     33 func AddRequestCompression(stack *middleware.Stack, disabled bool, minBytes int64, algorithms []string) error {
     34 	return stack.Serialize.Add(&requestCompression{
     35 		disableRequestCompression:   disabled,
     36 		requestMinCompressSizeBytes: minBytes,
     37 		compressAlgorithms:          algorithms,
     38 	}, middleware.After)
     39 }
     40 
     41 type requestCompression struct {
     42 	disableRequestCompression   bool
     43 	requestMinCompressSizeBytes int64
     44 	compressAlgorithms          []string
     45 }
     46 
     47 // ID returns the ID of the middleware
     48 func (m requestCompression) ID() string {
     49 	return "RequestCompression"
     50 }
     51 
     52 // HandleSerialize gzip compress the request's stream/body if enabled by config fields
     53 func (m requestCompression) HandleSerialize(
     54 	ctx context.Context, in middleware.SerializeInput, next middleware.SerializeHandler,
     55 ) (
     56 	out middleware.SerializeOutput, metadata middleware.Metadata, err error,
     57 ) {
     58 	if m.disableRequestCompression {
     59 		return next.HandleSerialize(ctx, in)
     60 	}
     61 	// still need to check requestMinCompressSizeBytes in case it is out of range after service client config
     62 	if m.requestMinCompressSizeBytes < 0 || m.requestMinCompressSizeBytes > MaxRequestMinCompressSizeBytes {
     63 		return out, metadata, fmt.Errorf("invalid range for min request compression size bytes %d, must be within 0 and 10485760 inclusively", m.requestMinCompressSizeBytes)
     64 	}
     65 
     66 	req, ok := in.Request.(*http.Request)
     67 	if !ok {
     68 		return out, metadata, fmt.Errorf("unknown request type %T", req)
     69 	}
     70 
     71 	for _, algorithm := range m.compressAlgorithms {
     72 		compressFunc := allowedAlgorithms[algorithm]
     73 		if compressFunc != nil {
     74 			if stream := req.GetStream(); stream != nil {
     75 				size, found, err := req.StreamLength()
     76 				if err != nil {
     77 					return out, metadata, fmt.Errorf("error while finding request stream length, %v", err)
     78 				} else if !found || size < m.requestMinCompressSizeBytes {
     79 					return next.HandleSerialize(ctx, in)
     80 				}
     81 
     82 				compressedBytes, err := compressFunc(stream)
     83 				if err != nil {
     84 					return out, metadata, fmt.Errorf("failed to compress request stream, %v", err)
     85 				}
     86 
     87 				var newReq *http.Request
     88 				if newReq, err = req.SetStream(bytes.NewReader(compressedBytes)); err != nil {
     89 					return out, metadata, fmt.Errorf("failed to set request stream, %v", err)
     90 				}
     91 				*req = *newReq
     92 
     93 				if val := req.Header.Get("Content-Encoding"); val != "" {
     94 					req.Header.Set("Content-Encoding", fmt.Sprintf("%s, %s", val, algorithm))
     95 				} else {
     96 					req.Header.Set("Content-Encoding", algorithm)
     97 				}
     98 			}
     99 			break
    100 		}
    101 	}
    102 
    103 	return next.HandleSerialize(ctx, in)
    104 }