From 5ddad4a5a2fc9671836dd6c4b27196b438fba2db Mon Sep 17 00:00:00 2001
From: Rick Burgess <richard.burgess@dazn.com>
Date: Tue, 17 Dec 2019 11:05:25 +0000
Subject: [PATCH] Handle throttle exceptions on public key send

---
 go.mod                      |  1 +
 go.sum                      |  2 ++
 main.go                     |  8 ++++++++
 pkg/sshutils/ec2endpoint.go | 13 +++++++++++--
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 056a5cb..626228d 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.13
 require (
 	github.com/aws/aws-sdk-go v1.25.48
 	github.com/blang/semver v3.5.1+incompatible
+	github.com/pkg/errors v0.8.1
 	github.com/rhysd/go-github-selfupdate v1.1.0
 	github.com/sirupsen/logrus v1.4.2
 	github.com/urfave/cli/v2 v2.0.0
diff --git a/go.sum b/go.sum
index bf3df8a..3286894 100644
--- a/go.sum
+++ b/go.sum
@@ -32,6 +32,8 @@ github.com/onsi/ginkgo v1.6.0 h1:Ix8l273rp3QzYgXSR+c8d1fTG7UPgYkOSELPhiY/YGw=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/gomega v1.4.2 h1:3mYCb7aPxS/RU7TI1y4rkEn1oKmPRjNJLNEXgw7MH2I=
 github.com/onsi/gomega v1.4.2/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
+github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rhysd/go-github-selfupdate v1.1.0 h1:+aMomy69YCYxJ6kr13nYIgAJWSB1kHK5M5YpbmjQkWo=
diff --git a/main.go b/main.go
index 817740f..0641f70 100644
--- a/main.go
+++ b/main.go
@@ -76,6 +76,10 @@ func main() {
 				Aliases:     []string{"lp"},
 				DefaultText: "local port to map to, defaults to tunnel port",
 			},
+			&cli.BoolFlag{
+				Name:        "debug",
+				DefaultText: "Print debug information",
+			},
 		},
 		Commands: []*cli.Command{
 			{
@@ -103,6 +107,10 @@ func SetupSignalHandlers() {
 }
 
 func run(c *cli.Context) error {
+	if c.Bool("debug") {
+		log.SetLevel(log.DebugLevel)
+	}
+
 	var tagName string
 	var tagValue string
 
diff --git a/pkg/sshutils/ec2endpoint.go b/pkg/sshutils/ec2endpoint.go
index 1e61f92..0f670b9 100644
--- a/pkg/sshutils/ec2endpoint.go
+++ b/pkg/sshutils/ec2endpoint.go
@@ -1,14 +1,15 @@
 package sshutils
 
 import (
-	"errors"
 	"fmt"
 	"strconv"
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	connect "github.com/aws/aws-sdk-go/service/ec2instanceconnect"
+	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
 )
@@ -94,8 +95,16 @@ func sendPublicKey(instance *ec2.Instance, user, publicKey string, client *conne
 		InstanceOSUser:   aws.String(user),
 		SSHPublicKey:     aws.String(publicKey),
 	})
+
 	if err != nil {
-		return err
+		if awsErr, ok := err.(awserr.Error); ok {
+			if awsErr.Code() == connect.ErrCodeThrottlingException {
+				log.Debug("Got throttling exception, usually just means the key is already valid")
+				return nil
+			}
+		}
+
+		return errors.Wrap(err, "send public key error")
 	}
 
 	if !*out.Success {