15061: API support for linking local account to remote user

[arvados.git] / lib / cloud / ec2 / ec2.go

diff --git a/lib/cloud/ec2/ec2.go b/lib/cloud/ec2/ec2.go
index b8a360e19fabd36777b15112e8f47ccaf3a7bab3..c5565d424559f0bba2841dd46df62d3af883cc19 100644 (file)
--- a/lib/cloud/ec2/ec2.go
+++ b/lib/cloud/ec2/ec2.go
@@ -5,11 +5,16 @@
 package ec2
 
 import (
+       "crypto/md5"
+       "crypto/rsa"
+       "crypto/sha1"
+       "crypto/x509"
        "encoding/base64"
        "encoding/json"
        "fmt"
-       "log"
+       "math/big"
        "strings"
+       "sync"
 
        "git.curoverse.com/arvados.git/lib/cloud"
        "git.curoverse.com/arvados.git/sdk/go/arvados"
@@ -34,10 +39,11 @@ type ec2InstanceSetConfig struct {
        SecurityGroupIDs []string
        SubnetID         string
        AdminUsername    string
-       KeyPairName      string
+       EBSVolumeType    string
 }
 
 type ec2Interface interface {
+       DescribeKeyPairs(input *ec2.DescribeKeyPairsInput) (*ec2.DescribeKeyPairsOutput, error)
        ImportKeyPair(input *ec2.ImportKeyPairInput) (*ec2.ImportKeyPairOutput, error)
        RunInstances(input *ec2.RunInstancesInput) (*ec2.Reservation, error)
        DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error)
@@ -50,7 +56,8 @@ type ec2InstanceSet struct {
        dispatcherID cloud.InstanceSetID
        logger       logrus.FieldLogger
        client       ec2Interface
-       importedKey  bool
+       keysMtx      sync.Mutex
+       keys         map[string]string
 }
 
 func newEC2InstanceSet(config json.RawMessage, dispatcherID cloud.InstanceSetID, logger logrus.FieldLogger) (prv cloud.InstanceSet, err error) {
@@ -69,9 +76,46 @@ func newEC2InstanceSet(config json.RawMessage, dispatcherID cloud.InstanceSetID,
                        "")).
                WithRegion(instanceSet.ec2config.Region)
        instanceSet.client = ec2.New(session.Must(session.NewSession(awsConfig)))
+       instanceSet.keys = make(map[string]string)
+       if instanceSet.ec2config.EBSVolumeType == "" {
+               instanceSet.ec2config.EBSVolumeType = "gp2"
+       }
        return instanceSet, nil
 }
 
+func awsKeyFingerprint(pk ssh.PublicKey) (md5fp string, sha1fp string, err error) {
+       // AWS key fingerprints don't use the usual key fingerprint
+       // you get from ssh-keygen or ssh.FingerprintLegacyMD5()
+       // (you can get that from md5.Sum(pk.Marshal())
+       //
+       // AWS uses the md5 or sha1 of the PKIX DER encoding of the
+       // public key, so calculate those fingerprints here.
+       var rsaPub struct {
+               Name string
+               E    *big.Int
+               N    *big.Int
+       }
+       if err := ssh.Unmarshal(pk.Marshal(), &rsaPub); err != nil {
+               return "", "", fmt.Errorf("agent: Unmarshal failed to parse public key: %v", err)
+       }
+       rsaPk := rsa.PublicKey{
+               E: int(rsaPub.E.Int64()),
+               N: rsaPub.N,
+       }
+       pkix, _ := x509.MarshalPKIXPublicKey(&rsaPk)
+       md5pkix := md5.Sum([]byte(pkix))
+       sha1pkix := sha1.Sum([]byte(pkix))
+       md5fp = ""
+       sha1fp = ""
+       for i := 0; i < len(md5pkix); i += 1 {
+               md5fp += fmt.Sprintf(":%02x", md5pkix[i])
+       }
+       for i := 0; i < len(sha1pkix); i += 1 {
+               sha1fp += fmt.Sprintf(":%02x", sha1pkix[i])
+       }
+       return md5fp[1:], sha1fp[1:], nil
+}
+
 func (instanceSet *ec2InstanceSet) Create(
        instanceType arvados.InstanceType,
        imageID cloud.ImageID,
@@ -79,13 +123,39 @@ func (instanceSet *ec2InstanceSet) Create(
        initCommand cloud.InitCommand,
        publicKey ssh.PublicKey) (cloud.Instance, error) {
 
-       if !instanceSet.importedKey {
-               instanceSet.client.ImportKeyPair(&ec2.ImportKeyPairInput{
-                       KeyName:           &instanceSet.ec2config.KeyPairName,
-                       PublicKeyMaterial: ssh.MarshalAuthorizedKey(publicKey),
+       md5keyFingerprint, sha1keyFingerprint, err := awsKeyFingerprint(publicKey)
+       if err != nil {
+               return nil, fmt.Errorf("Could not make key fingerprint: %v", err)
+       }
+       instanceSet.keysMtx.Lock()
+       var keyname string
+       var ok bool
+       if keyname, ok = instanceSet.keys[md5keyFingerprint]; !ok {
+               keyout, err := instanceSet.client.DescribeKeyPairs(&ec2.DescribeKeyPairsInput{
+                       Filters: []*ec2.Filter{&ec2.Filter{
+                               Name:   aws.String("fingerprint"),
+                               Values: []*string{&md5keyFingerprint, &sha1keyFingerprint},
+                       }},
                })
-               instanceSet.importedKey = true
+               if err != nil {
+                       return nil, fmt.Errorf("Could not search for keypair: %v", err)
+               }
+
+               if len(keyout.KeyPairs) > 0 {
+                       keyname = *(keyout.KeyPairs[0].KeyName)
+               } else {
+                       keyname = "arvados-dispatch-keypair-" + md5keyFingerprint
+                       _, err := instanceSet.client.ImportKeyPair(&ec2.ImportKeyPairInput{
+                               KeyName:           &keyname,
+                               PublicKeyMaterial: ssh.MarshalAuthorizedKey(publicKey),
+                       })
+                       if err != nil {
+                               return nil, fmt.Errorf("Could not import keypair: %v", err)
+                       }
+               }
+               instanceSet.keys[md5keyFingerprint] = keyname
        }
+       instanceSet.keysMtx.Unlock()
 
        ec2tags := []*ec2.Tag{
                &ec2.Tag{
@@ -109,7 +179,7 @@ func (instanceSet *ec2InstanceSet) Create(
                InstanceType: &instanceType.ProviderType,
                MaxCount:     aws.Int64(1),
                MinCount:     aws.Int64(1),
-               KeyName:      &instanceSet.ec2config.KeyPairName,
+               KeyName:      &keyname,
 
                NetworkInterfaces: []*ec2.InstanceNetworkInterfaceSpecification{
                        &ec2.InstanceNetworkInterfaceSpecification{
@@ -129,13 +199,13 @@ func (instanceSet *ec2InstanceSet) Create(
                        }},
        }
 
-       if instanceType.AttachScratch {
+       if instanceType.AddedScratch > 0 {
                rii.BlockDeviceMappings = []*ec2.BlockDeviceMapping{&ec2.BlockDeviceMapping{
                        DeviceName: aws.String("/dev/xvdt"),
                        Ebs: &ec2.EbsBlockDevice{
                                DeleteOnTermination: aws.Bool(true),
-                               VolumeSize:          aws.Int64((int64(instanceType.Scratch) / 1000000000) + 1),
-                               VolumeType:          aws.String("gp2"),
+                               VolumeSize:          aws.Int64((int64(instanceType.AddedScratch) + (1<<30 - 1)) >> 30),
+                               VolumeType:          &instanceSet.ec2config.EBSVolumeType,
                        }}}
        }
 
@@ -175,7 +245,9 @@ func (instanceSet *ec2InstanceSet) Instances(cloud.InstanceTags) (instances []cl
 
                for _, rsv := range dio.Reservations {
                        for _, inst := range rsv.Instances {
-                               instances = append(instances, &ec2Instance{instanceSet, inst})
+                               if *inst.State.Name != "shutting-down" && *inst.State.Name != "terminated" {
+                                       instances = append(instances, &ec2Instance{instanceSet, inst})
+                               }
                        }
                }
                if dio.NextToken == nil {
@@ -240,7 +312,6 @@ func (inst *ec2Instance) Tags() cloud.InstanceTags {
 }
 
 func (inst *ec2Instance) Destroy() error {
-       log.Printf("terminating %v", *inst.instance.InstanceId)
        _, err := inst.provider.client.TerminateInstances(&ec2.TerminateInstancesInput{
                InstanceIds: []*string{inst.instance.InstanceId},
        })