golang.org/x/build@v0.0.0-20240506185731-218518f32b70/internal/cloud/fake_aws.go

// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cloud

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	mrand "math/rand"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/service/ec2"
)

func init() { mrand.Seed(time.Now().UnixNano()) }

// FakeAWSClient provides a fake AWS Client used to test the AWS client
// functionality.
type FakeAWSClient struct {
	mu            sync.RWMutex
	instances     map[string]*Instance
	instanceTypes []*InstanceType
	serviceQuotas map[serviceQuotaKey]int64
}

// serviceQuotaKey should be used as the key in the serviceQuotas map.
type serviceQuotaKey struct {
	code    string
	service string
}

// NewFakeAWSClient crates a fake AWS client.
func NewFakeAWSClient() *FakeAWSClient {
	return &FakeAWSClient{
		instances: make(map[string]*Instance),
		instanceTypes: []*InstanceType{
			&InstanceType{"ab.large", 10},
			&InstanceType{"ab.xlarge", 20},
			&InstanceType{"ab.small", 30},
		},
		serviceQuotas: map[serviceQuotaKey]int64{
			serviceQuotaKey{QuotaCodeCPUOnDemand, QuotaServiceEC2}: 384,
		},
	}
}

// Instance returns the `Instance` record for the requested instance. The instance record will
// return records for recently terminated instances. If an instance is not found an error will
// be returned.
func (f *FakeAWSClient) Instance(ctx context.Context, instID string) (*Instance, error) {
	if ctx == nil || instID == "" {
		return nil, errors.New("invalid params")
	}
	f.mu.RLock()
	defer f.mu.RUnlock()

	inst, ok := f.instances[instID]
	if !ok {
		return nil, errors.New("instance not found")
	}
	return copyInstance(inst), nil
}

// Instances retrieves all EC2 instances in a region which have not been terminated or stopped.
func (f *FakeAWSClient) RunningInstances(ctx context.Context) ([]*Instance, error) {
	if ctx == nil {
		return nil, errors.New("invalid params")
	}
	f.mu.RLock()
	defer f.mu.RUnlock()

	instances := make([]*Instance, 0, len(f.instances))
	for _, inst := range f.instances {
		if inst.State != ec2.InstanceStateNameRunning && inst.State != ec2.InstanceStateNamePending {
			continue
		}
		instances = append(instances, copyInstance(inst))
	}
	return instances, nil
}

// InstanceTypesARM retrieves all EC2 instance types in a region which support the ARM64 architecture.
func (f *FakeAWSClient) InstanceTypesARM(ctx context.Context) ([]*InstanceType, error) {
	if ctx == nil {
		return nil, errors.New("invalid params")
	}
	f.mu.RLock()
	defer f.mu.RUnlock()

	instanceTypes := make([]*InstanceType, 0, len(f.instanceTypes))
	for _, it := range f.instanceTypes {
		instanceTypes = append(instanceTypes, &InstanceType{it.Type, it.CPU})
	}
	return instanceTypes, nil
}

// Quota retrieves the requested service quota for the service.
func (f *FakeAWSClient) Quota(ctx context.Context, service, code string) (int64, error) {
	if ctx == nil || service == "" || code == "" {
		return 0, errors.New("invalid params")
	}
	f.mu.RLock()
	defer f.mu.RUnlock()

	v, ok := f.serviceQuotas[serviceQuotaKey{code, service}]
	if !ok {
		return 0, errors.New("service quota not found")
	}
	return v, nil
}

// CreateInstance creates an EC2 VM instance.
func (f *FakeAWSClient) CreateInstance(ctx context.Context, config *EC2VMConfiguration) (*Instance, error) {
	if ctx == nil || config == nil {
		return nil, errors.New("invalid params")
	}
	if config.ImageID == "" {
		return nil, errors.New("invalid Image ID")
	}
	if config.Type == "" {
		return nil, errors.New("invalid Type")
	}
	if config.Zone == "" {
		return nil, errors.New("invalid Zone")
	}
	f.mu.Lock()
	defer f.mu.Unlock()

	inst := &Instance{
		CPUCount:          4,
		CreatedAt:         time.Now(),
		Description:       config.Description,
		ID:                fmt.Sprintf("instance-%s", randHex(10)),
		IPAddressExternal: randIPv4(),
		IPAddressInternal: randIPv4(),
		ImageID:           config.ImageID,
		Name:              config.Name,
		SSHKeyID:          config.SSHKeyID,
		SecurityGroups:    config.SecurityGroups,
		State:             ec2.InstanceStateNameRunning,
		Tags:              make(map[string]string),
		Type:              config.Type,
		Zone:              config.Zone,
	}
	for k, v := range config.Tags {
		inst.Tags[k] = v
	}
	f.instances[inst.ID] = inst
	return copyInstance(inst), nil
}

// DestroyInstances terminates EC2 VM instances.
func (f *FakeAWSClient) DestroyInstances(ctx context.Context, instIDs ...string) error {
	if ctx == nil || len(instIDs) == 0 {
		return errors.New("invalid params")
	}
	f.mu.Lock()
	defer f.mu.Unlock()

	for _, id := range instIDs {
		inst, ok := f.instances[id]
		if !ok {
			return errors.New("instance not found")
		}
		inst.State = ec2.InstanceStateNameTerminated
	}
	return nil
}

// WaitUntilInstanceRunning returns when an instance has transitioned into the running state.
func (f *FakeAWSClient) WaitUntilInstanceRunning(ctx context.Context, instID string) error {
	if ctx == nil || instID == "" {
		return errors.New("invalid params")
	}
	f.mu.RLock()
	defer f.mu.RUnlock()

	inst, ok := f.instances[instID]
	if !ok {
		return errors.New("instance not found")
	}
	if inst.State != ec2.InstanceStateNameRunning {
		return errors.New("timed out waiting for instance to enter running state")
	}
	return nil
}

// copyInstance copies the contents of a pointer to an instance and returns a newly created
// instance with the same data as the original instance.
func copyInstance(inst *Instance) *Instance {
	i := &Instance{
		CPUCount:          inst.CPUCount,
		CreatedAt:         inst.CreatedAt,
		Description:       inst.Description,
		ID:                inst.ID,
		IPAddressExternal: inst.IPAddressExternal,
		IPAddressInternal: inst.IPAddressInternal,
		ImageID:           inst.ImageID,
		Name:              inst.Name,
		SSHKeyID:          inst.SSHKeyID,
		SecurityGroups:    inst.SecurityGroups,
		State:             inst.State,
		Tags:              make(map[string]string),
		Type:              inst.Type,
		Zone:              inst.Zone,
	}
	for k, v := range inst.Tags {
		i.Tags[k] = v
	}
	return i
}

// randHex creates a random hex string of length n.
func randHex(n int) string {
	buf := make([]byte, n/2+1)
	_, _ = rand.Read(buf)
	return fmt.Sprintf("%x", buf)[:n]
}

// randIPv4 creates a random IPv4 address.
func randIPv4() string {
	return fmt.Sprintf("%d.%d.%d.%d", mrand.Intn(255), mrand.Intn(255), mrand.Intn(255), mrand.Intn(255))
}