Merge pull request #20 from rhnvrm/feat-iam-isolate

feat: add support for timeout and customization in IAM setup.

README.md

@@ -2,8 +2,8 @@
 
 ## Overview [![GoDoc](https://godoc.org/github.com/rhnvrm/simples3?status.svg)](https://godoc.org/github.com/rhnvrm/simples3) [![Go Report Card](https://goreportcard.com/badge/github.com/rhnvrm/simples3)](https://goreportcard.com/report/github.com/rhnvrm/simples3) [![GoCover](https://gocover.io/_badge/github.com/rhnvrm/simples3)](https://gocover.io/_badge/github.com/rhnvrm/simples3) [![Zerodha Tech](https://zerodha.tech/static/images/github-badge.svg)](https://zerodha.tech) 
 
-SimpleS3 is a golang library for uploading and deleting objects 
-on S3 buckets using REST API calls or Presigned URLs signed 
+SimpleS3 is a Go library for manipulating objects 
+in S3 buckets using REST API calls or Presigned URLs signed 
 using AWS Signature Version 4.
 
 ## Install
@@ -54,6 +54,7 @@ file, _ := s3.FileDownload(simples3.DownloadInput{
     Bucket:    AWSBucket,
     ObjectKey: "test.txt",
 })
+
 data, _ := ioutil.ReadAll(file)
 file.Close()
 

simples3.go

@@ -24,8 +24,7 @@ import (
 
 const (
 	securityCredentialsURL = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
-)
-const (
+
 	// AMZMetaPrefix to prefix metadata key.
 	AMZMetaPrefix = "x-amz-meta-"
 )
@@ -143,54 +142,79 @@ func New(region, accessKey, secretKey string) *S3 {
 // NewUsingIAM automatically generates an Instance of S3
 // using instance metatdata.
 func NewUsingIAM(region string) (*S3, error) {
-	return newUsingIAMImpl(securityCredentialsURL, region)
+	return newUsingIAM(
+		&http.Client{
+			// Set a timeout of 3 seconds for AWS IAM Calls.
+			Timeout: time.Second * 3, //nolint:gomnd
+		}, securityCredentialsURL, region)
 }
 
-func newUsingIAMImpl(baseURL, region string) (*S3, error) {
-	// Get the IAM role
-	resp, err := http.Get(baseURL)
+// fetchIAMData fetches the IAM data from the given URL.
+// In case of a normal AWS setup, baseURL would be securityCredentialsURL.
+// You can use this method, to manually fetch IAM data from a custom
+// endpoint and pass it to SetIAMData.
+func fetchIAMData(cl *http.Client, baseURL string) (IAMResponse, error) {
+	resp, err := cl.Get(baseURL)
 	if err != nil {
-		return nil, err
+		return IAMResponse{}, err
 	}
 	defer resp.Body.Close()
+
 	if resp.StatusCode != http.StatusOK {
-		return nil, errors.New(http.StatusText(resp.StatusCode))
+		return IAMResponse{}, fmt.Errorf("Error fetching IAM data: %s", resp.Status)
 	}
 
 	role, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return IAMResponse{}, err
 	}
 
 	resp, err = http.Get(baseURL + "/" + string(role))
 	if err != nil {
-		return nil, err
+		return IAMResponse{}, err
 	}
 	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
-		return nil, errors.New(http.StatusText(resp.StatusCode))
+		return IAMResponse{}, errors.New(http.StatusText(resp.StatusCode))
 	}
 
-	var jsonResp IAMResponse
+	var jResp IAMResponse
 	jsonString, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return IAMResponse{}, err
 	}
 
-	if err := json.Unmarshal(jsonString, &jsonResp); err != nil {
-		return nil, err
+	if err := json.Unmarshal(jsonString, &jResp); err != nil {
+		return IAMResponse{}, err
+	}
+
+	return jResp, nil
+}
+
+func newUsingIAM(cl *http.Client, credUrl, region string) (*S3, error) {
+	// Get the IAM role
+	iamResp, err := fetchIAMData(cl, credUrl)
+	if err != nil {
+		return nil, fmt.Errorf("Error fetching IAM data: %w", err)
 	}
 
 	return &S3{
 		Region:    region,
-		AccessKey: jsonResp.AccessKeyID,
-		SecretKey: jsonResp.SecretAccessKey,
-		Token:     jsonResp.Token,
+		AccessKey: iamResp.AccessKeyID,
+		SecretKey: iamResp.SecretAccessKey,
+		Token:     iamResp.Token,
 
 		URIFormat: "https://s3.%s.amazonaws.com/%s",
 	}, nil
 }
 
+// setIAMData sets the IAM data on the S3 instance.
+func (s3 *S3) setIAMData(iamResp IAMResponse) {
+	s3.AccessKey = iamResp.AccessKeyID
+	s3.SecretKey = iamResp.SecretAccessKey
+	s3.Token = iamResp.Token
+}
+
 func (s3 *S3) getClient() *http.Client {
 	if s3.Client == nil {
 		return http.DefaultClient

simples3_test.go

@@ -2,12 +2,15 @@ package simples3
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"testing"
+	"time"
 )
 
 type tConfig struct {
@@ -398,6 +401,12 @@ func TestS3_NewUsingIAM(t *testing.T) {
 				"SecretAccessKey" : "abc","Token" : "abc",
 				"Expiration" : "2018-12-24T16:24:59Z"}`
 	)
+
+	tsFail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(2 * time.Second)
+	}))
+	defer tsFail.Close()
+
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Method != "GET" {
 			t.Errorf("Expected 'GET' request, got '%s'", r.Method)
@@ -414,10 +423,23 @@ func TestS3_NewUsingIAM(t *testing.T) {
 	}))
 	defer ts.Close()
 
-	s3, err := newUsingIAMImpl(ts.URL, "abc")
+	// Test for timeout.
+	_, err := newUsingIAM(&http.Client{Timeout: 1 * time.Second}, tsFail.URL, "abc")
+	if err == nil {
+		t.Errorf("Expected error, got nil")
+	} else {
+		var timeoutError net.Error
+		if errors.As(err, &timeoutError); !timeoutError.Timeout() {
+			t.Errorf("newUsingIAM() timeout check. got error = %v", err)
+		}
+	}
+
+	// Test for successful IAM fetch.
+	s3, err := newUsingIAM(http.DefaultClient, ts.URL, "abc")
 	if err != nil {
-		t.Errorf("S3.FileDelete() error = %v", err)
+		t.Errorf("newUsingIAM() error = %v", err)
 	}
+
 	if s3.AccessKey != "abc" && s3.SecretKey != "abc" && s3.Region != "abc" {
 		t.Errorf("S3.FileDelete() got = %v", s3)
 	}