Update option handling in EC2Client and add custom HTTP client (#17)

plugin/option.go

@@ -0,0 +1,35 @@
+package plugin
+
+// getOpts iterates the inbound Options and returns a struct
+func getOpts(opt ...Option) (options, error) {
+	opts := getDefaultOptions()
+	for _, o := range opt {
+		if o == nil {
+			continue
+		}
+		if err := o(&opts); err != nil {
+			return options{}, err
+		}
+	}
+	return opts, nil
+}
+
+// Option - how Options are passed as arguments
+type Option func(*options) error
+
+// options = how options are represented
+type options struct {
+	withRegion string
+}
+
+func getDefaultOptions() options {
+	return options{}
+}
+
+// WithRegion contains the region to use
+func WithRegion(with string) Option {
+	return func(o *options) error {
+		o.withRegion = with
+		return nil
+	}
+}

plugin/plugin.go

@@ -272,7 +272,7 @@ func (p *AwsPlugin) OnCreateSet(ctx context.Context, req *pb.OnCreateSetRequest)
 		return nil, err
 	}
 
-	ec2Client, err := state.EC2Client(catalogAttributes.Region)
+	ec2Client, err := state.EC2Client(WithRegion(catalogAttributes.Region))
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error getting EC2 client: %s", err)
 	}
@@ -335,7 +335,7 @@ func (p *AwsPlugin) OnUpdateSet(ctx context.Context, req *pb.OnUpdateSetRequest)
 		return nil, err
 	}
 
-	ec2Client, err := state.EC2Client(catalogAttributes.Region)
+	ec2Client, err := state.EC2Client(WithRegion(catalogAttributes.Region))
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error getting EC2 client: %s", err)
 	}
@@ -427,7 +427,7 @@ func (p *AwsPlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (*p
 		}
 	}
 
-	ec2Client, err := state.EC2Client(catalogAttributes.Region)
+	ec2Client, err := state.EC2Client(WithRegion(catalogAttributes.Region))
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error getting EC2 client: %s", err)
 	}

plugin/state.go

@@ -3,6 +3,8 @@ package plugin
 import (
 	"errors"
 	"fmt"
+	"net"
+	"net/http"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -259,17 +261,40 @@ func (s *awsCatalogPersistedState) GetSession() (*session.Session, error) {
 	return c.GetSession()
 }
 
+var customClient = &http.Client{
+	Transport: &http.Transport{
+		Proxy: http.ProxyFromEnvironment,
+		DialContext: (&net.Dialer{
+			Timeout:   30 * time.Second,
+			KeepAlive: 30 * time.Second,
+			DualStack: true,
+		}).DialContext,
+	},
+}
+
 // EC2Client returns a configured EC2 client based on the session
 // information stored in the state.
-func (s *awsCatalogPersistedState) EC2Client(region string) (ec2iface.EC2API, error) {
+func (s *awsCatalogPersistedState) EC2Client(opt ...Option) (ec2iface.EC2API, error) {
+	opts, err := getOpts(opt...)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing options when fetching EC2 client: %w", err)
+	}
+
 	sess, err := s.GetSession()
 	if err != nil {
-		return nil, fmt.Errorf("error getting AWS session: %w", err)
+		return nil, fmt.Errorf("error getting AWS session when fetching EC2 client: %w", err)
 	}
 
 	if s.testEC2APIFunc != nil {
-		return s.testEC2APIFunc(sess, aws.NewConfig().WithRegion(region))
+		return s.testEC2APIFunc(sess, aws.NewConfig().WithRegion(opts.withRegion))
 	}
 
-	return ec2.New(sess, aws.NewConfig().WithRegion(region)), nil
+	cfg := aws.NewConfig()
+	if opts.withRegion != "" {
+		cfg = cfg.WithRegion(opts.withRegion)
+	}
+
+	cfg.HTTPClient = customClient
+
+	return ec2.New(sess, cfg), nil
 }