util.go

Documentation: github.com/aws/aws-sdk-go-v2/internal/awstesting

     1  package awstesting
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net/http"
     7  	"os"
     8  	"runtime"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/aws/aws-sdk-go-v2/aws"
    13  )
    14  
    15  // ZeroReader is a io.Reader which will always write zeros to the byte slice provided.
    16  type ZeroReader struct{}
    17  
    18  // Read fills the provided byte slice with zeros returning the number of bytes written.
    19  func (r *ZeroReader) Read(b []byte) (int, error) {
    20  	for i := 0; i < len(b); i++ {
    21  		b[i] = 0
    22  	}
    23  	return len(b), nil
    24  }
    25  
    26  // ReadCloser is a io.ReadCloser for unit testing.
    27  // Designed to test for leaks and whether a handle has
    28  // been closed
    29  type ReadCloser struct {
    30  	Size     int
    31  	Closed   bool
    32  	set      bool
    33  	FillData func(bool, []byte, int, int)
    34  }
    35  
    36  // Read will call FillData and fill it with whatever data needed.
    37  // Decrements the size until zero, then return io.EOF.
    38  func (r *ReadCloser) Read(b []byte) (int, error) {
    39  	if r.Closed {
    40  		return 0, io.EOF
    41  	}
    42  
    43  	delta := len(b)
    44  	if delta > r.Size {
    45  		delta = r.Size
    46  	}
    47  	r.Size -= delta
    48  
    49  	for i := 0; i < delta; i++ {
    50  		b[i] = 'a'
    51  	}
    52  
    53  	if r.FillData != nil {
    54  		r.FillData(r.set, b, r.Size, delta)
    55  	}
    56  	r.set = true
    57  
    58  	if r.Size > 0 {
    59  		return delta, nil
    60  	}
    61  	return delta, io.EOF
    62  }
    63  
    64  // Close sets Closed to true and returns no error
    65  func (r *ReadCloser) Close() error {
    66  	r.Closed = true
    67  	return nil
    68  }
    69  
    70  // A FakeContext provides a simple stub implementation of a Context
    71  type FakeContext struct {
    72  	Error  error
    73  	DoneCh chan struct{}
    74  }
    75  
    76  // Deadline always will return not set
    77  func (c *FakeContext) Deadline() (deadline time.Time, ok bool) {
    78  	return time.Time{}, false
    79  }
    80  
    81  // Done returns a read channel for listening to the Done event
    82  func (c *FakeContext) Done() <-chan struct{} {
    83  	return c.DoneCh
    84  }
    85  
    86  // Err returns the error, is nil if not set.
    87  func (c *FakeContext) Err() error {
    88  	return c.Error
    89  }
    90  
    91  // Value ignores the Value and always returns nil
    92  func (c *FakeContext) Value(key interface{}) interface{} {
    93  	return nil
    94  }
    95  
    96  // StashEnv stashes the current environment variables except variables listed in envToKeepx
    97  // Returns an function to pop out old environment
    98  func StashEnv(envToKeep ...string) []string {
    99  	if runtime.GOOS == "windows" {
   100  		envToKeep = append(envToKeep, "ComSpec")
   101  		envToKeep = append(envToKeep, "SYSTEM32")
   102  		envToKeep = append(envToKeep, "SYSTEMROOT")
   103  	}
   104  	envToKeep = append(envToKeep, "PATH", "HOME", "USERPROFILE")
   105  	extraEnv := getEnvs(envToKeep)
   106  	originalEnv := os.Environ()
   107  	os.Clearenv() // clear env
   108  	for key, val := range extraEnv {
   109  		os.Setenv(key, val)
   110  	}
   111  	return originalEnv
   112  }
   113  
   114  // PopEnv takes the list of the environment values and injects them into the
   115  // process's environment variable data. Clears any existing environment values
   116  // that may already exist.
   117  func PopEnv(env []string) {
   118  	os.Clearenv()
   119  
   120  	for _, e := range env {
   121  		p := strings.SplitN(e, "=", 2)
   122  		k, v := p[0], ""
   123  		if len(p) > 1 {
   124  			v = p[1]
   125  		}
   126  		os.Setenv(k, v)
   127  	}
   128  }
   129  
   130  // MockCredentialsProvider is a type that can be used to mock out credentials
   131  // providers
   132  type MockCredentialsProvider struct {
   133  	RetrieveFn   func(ctx context.Context) (aws.Credentials, error)
   134  	InvalidateFn func()
   135  }
   136  
   137  // Retrieve calls the RetrieveFn
   138  func (p MockCredentialsProvider) Retrieve(ctx context.Context) (aws.Credentials, error) {
   139  	return p.RetrieveFn(ctx)
   140  }
   141  
   142  // Invalidate calls the InvalidateFn
   143  func (p MockCredentialsProvider) Invalidate() {
   144  	p.InvalidateFn()
   145  }
   146  
   147  func getEnvs(envs []string) map[string]string {
   148  	extraEnvs := make(map[string]string)
   149  	for _, env := range envs {
   150  		if val, ok := os.LookupEnv(env); ok && len(val) > 0 {
   151  			extraEnvs[env] = val
   152  		}
   153  	}
   154  	return extraEnvs
   155  }
   156  
   157  const (
   158  	signaturePreambleSigV4  = "AWS4-HMAC-SHA256"
   159  	signaturePreambleSigV4A = "AWS4-ECDSA-P256-SHA256"
   160  )
   161  
   162  // SigV4Signature represents a parsed sigv4 or sigv4a signature.
   163  type SigV4Signature struct {
   164  	Preamble      string   // e.g. AWS4-HMAC-SHA256, AWS4-ECDSA-P256-SHA256
   165  	SigningName   string   // generally the service name e.g. "s3"
   166  	SigningRegion string   // for sigv4a this is the region-set header as-is
   167  	SignedHeaders []string // list of signed headers
   168  	Signature     string   // calculated signature
   169  }
   170  
   171  // ParseSigV4Signature deconstructs a sigv4 or sigv4a signature from a set of
   172  // request headers.
   173  func ParseSigV4Signature(header http.Header) *SigV4Signature {
   174  	auth := header.Get("Authorization")
   175  
   176  	preamble, after, _ := strings.Cut(auth, " ")
   177  	credential, after, _ := strings.Cut(after, ", ")
   178  	signedHeaders, signature, _ := strings.Cut(after, ", ")
   179  
   180  	credentialParts := strings.Split(credential, "/")
   181  
   182  	// sigv4  : AccessKeyID/DateString/SigningRegion/SigningName/SignatureID
   183  	// sigv4a : AccessKeyID/DateString/SigningName/SignatureID, region set on
   184  	//          header
   185  	var signingName, signingRegion string
   186  	if preamble == signaturePreambleSigV4 {
   187  		signingName = credentialParts[3]
   188  		signingRegion = credentialParts[2]
   189  	} else if preamble == signaturePreambleSigV4A {
   190  		signingName = credentialParts[2]
   191  		signingRegion = header.Get("X-Amz-Region-Set")
   192  	}
   193  
   194  	return &SigV4Signature{
   195  		Preamble:      preamble,
   196  		SigningName:   signingName,
   197  		SigningRegion: signingRegion,
   198  		SignedHeaders: strings.Split(signedHeaders, ";"),
   199  		Signature:     signature,
   200  	}
   201  }
   202
View as plain text