github.com/aavshr/aws-sdk-go@v1.41.3/service/dynamodb/customizations.go (about)

     1  package dynamodb
     2  
     3  import (
     4  	"bytes"
     5  	"hash/crc32"
     6  	"io"
     7  	"io/ioutil"
     8  	"strconv"
     9  	"time"
    10  
    11  	"github.com/aavshr/aws-sdk-go/aws"
    12  	"github.com/aavshr/aws-sdk-go/aws/awserr"
    13  	"github.com/aavshr/aws-sdk-go/aws/client"
    14  	"github.com/aavshr/aws-sdk-go/aws/request"
    15  )
    16  
    17  func init() {
    18  	initClient = func(c *client.Client) {
    19  		if c.Config.Retryer == nil {
    20  			// Only override the retryer with a custom one if the config
    21  			// does not already contain a retryer
    22  			setCustomRetryer(c)
    23  		}
    24  
    25  		c.Handlers.Build.PushBack(disableCompression)
    26  		c.Handlers.Unmarshal.PushFront(validateCRC32)
    27  	}
    28  }
    29  
    30  func setCustomRetryer(c *client.Client) {
    31  	maxRetries := aws.IntValue(c.Config.MaxRetries)
    32  	if c.Config.MaxRetries == nil || maxRetries == aws.UseServiceDefaultRetries {
    33  		maxRetries = 10
    34  	}
    35  
    36  	c.Retryer = client.DefaultRetryer{
    37  		NumMaxRetries: maxRetries,
    38  		MinRetryDelay: 50 * time.Millisecond,
    39  	}
    40  }
    41  
    42  func drainBody(b io.ReadCloser, length int64) (out *bytes.Buffer, err error) {
    43  	if length < 0 {
    44  		length = 0
    45  	}
    46  	buf := bytes.NewBuffer(make([]byte, 0, length))
    47  
    48  	if _, err = buf.ReadFrom(b); err != nil {
    49  		return nil, err
    50  	}
    51  	if err = b.Close(); err != nil {
    52  		return nil, err
    53  	}
    54  	return buf, nil
    55  }
    56  
    57  func disableCompression(r *request.Request) {
    58  	r.HTTPRequest.Header.Set("Accept-Encoding", "identity")
    59  }
    60  
    61  func validateCRC32(r *request.Request) {
    62  	if r.Error != nil {
    63  		return // already have an error, no need to verify CRC
    64  	}
    65  
    66  	// Checksum validation is off, skip
    67  	if aws.BoolValue(r.Config.DisableComputeChecksums) {
    68  		return
    69  	}
    70  
    71  	// Try to get CRC from response
    72  	header := r.HTTPResponse.Header.Get("X-Amz-Crc32")
    73  	if header == "" {
    74  		return // No header, skip
    75  	}
    76  
    77  	expected, err := strconv.ParseUint(header, 10, 32)
    78  	if err != nil {
    79  		return // Could not determine CRC value, skip
    80  	}
    81  
    82  	buf, err := drainBody(r.HTTPResponse.Body, r.HTTPResponse.ContentLength)
    83  	if err != nil { // failed to read the response body, skip
    84  		return
    85  	}
    86  
    87  	// Reset body for subsequent reads
    88  	r.HTTPResponse.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
    89  
    90  	// Compute the CRC checksum
    91  	crc := crc32.ChecksumIEEE(buf.Bytes())
    92  
    93  	if crc != uint32(expected) {
    94  		// CRC does not match, set a retryable error
    95  		r.Retryable = aws.Bool(true)
    96  		r.Error = awserr.New("CRC32CheckFailed", "CRC32 integrity check failed", nil)
    97  	}
    98  }