diff --git a/plugin/option.go b/plugin/option.go new file mode 100644 index 0000000..2a8e3e5 --- /dev/null +++ b/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 + } +} diff --git a/plugin/plugin.go b/plugin/plugin.go index 48998dd..000eaa7 100644 --- a/plugin/plugin.go +++ b/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) } diff --git a/plugin/state.go b/plugin/state.go index 982abef..75dca2a 100644 --- a/plugin/state.go +++ b/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 }