From 32210f52b024c113d4a5d08444b0aa24e5d85b33 Mon Sep 17 00:00:00 2001
From: Radon Rosborough <radon.neon@gmail.com>
Date: Mon, 5 Jul 2021 00:05:32 +0000
Subject: [PATCH] Awesome

---
 supervisor/go.mod      |  1 +
 supervisor/go.sum      |  2 ++
 supervisor/src/main.go | 45 ++++++++++++++++++++++++++++++++++++++++++
 tf/iam.tf              | 21 ++++++++++++++++++++
 4 files changed, 69 insertions(+)

diff --git a/supervisor/go.mod b/supervisor/go.mod
index 116938f..3021b2a 100644
--- a/supervisor/go.mod
+++ b/supervisor/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.7.0 // indirect
 	github.com/aws/aws-sdk-go-v2/config v1.4.1 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.3.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ecr v1.4.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.11.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.5.0 // indirect
 	github.com/caarlos0/env/v6 v6.6.2 // indirect
diff --git a/supervisor/go.sum b/supervisor/go.sum
index ed79512..5bd5fe1 100644
--- a/supervisor/go.sum
+++ b/supervisor/go.sum
@@ -10,6 +10,8 @@ github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.3.1 h1:ag1MjvYmE8hnvl2/3LYOog
 github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.3.1/go.mod h1:WXrj1wxGcYFfQ6H4xqsbVziISWQT55SlpX8B5+EqLOw=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.1.0 h1:DJq/vXXF+LAFaa/kQX9C6arlf4xX4uaaqGWIyAKOCpM=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.1.0/go.mod h1:qGQ/9IfkZonRNSNLE99/yBJ7EPA/h8jlWEqtJCcaj+Q=
+github.com/aws/aws-sdk-go-v2/service/ecr v1.4.0 h1:cgMcR4Y2JFhWHFDNiVYLApc5kSaGK0geqqL/2XvP77M=
+github.com/aws/aws-sdk-go-v2/service/ecr v1.4.0/go.mod h1:66eKvbrtxgZWfVHNwdncN8vciDvc00gX2flcATKqLYQ=
 github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.2.0 h1:wfI4yrOCMAGdHaEreQ65ycSmPLVc2Q82O+r7ZxYTynA=
 github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.2.0/go.mod h1:2Kc2Pybp1Hr2ZCCOz78mWnNSZYEKKBQgNcizVGk9sko=
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.2.0 h1:g2npzssI/6XsoQaPYCxliMFeC5iNKKvO0aC+/wWOE0A=
diff --git a/supervisor/src/main.go b/supervisor/src/main.go
index 42bd8a6..33bf915 100644
--- a/supervisor/src/main.go
+++ b/supervisor/src/main.go
@@ -1,7 +1,9 @@
 package main
 
 import (
+	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -21,6 +23,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	awsConfig "github.com/aws/aws-sdk-go-v2/config"
 	s3manager "github.com/aws/aws-sdk-go-v2/feature/s3/manager"
+	"github.com/aws/aws-sdk-go-v2/service/ecr"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/caarlos0/env/v6"
@@ -59,6 +62,7 @@ type supervisor struct {
 	awsAccountNumber string
 	awsRegion string
 	s3 *s3.Client
+	ecr *ecr.Client
 
 	reloadLock sync.Mutex
 	reloadInProgress bool
@@ -204,6 +208,46 @@ func (sv *supervisor) reloadWithScheduling() {
 var rijuImageRegexp = regexp.MustCompile(`(?:^|/)riju:([^<>]+)$`)
 
 func (sv *supervisor) reload() error {
+	sv.status("getting access token from ECR")
+	ecrResp, err := sv.ecr.GetAuthorizationToken(
+		context.Background(),
+		&ecr.GetAuthorizationTokenInput{},
+	)
+	if err != nil {
+		return err
+	}
+	if len(ecrResp.AuthorizationData) != 1 {
+		return fmt.Errorf(
+			"got unexpected number (%d) of authorization tokens",
+			len(ecrResp.AuthorizationData),
+		)
+	}
+	authInfo, err := base64.StdEncoding.DecodeString(*ecrResp.AuthorizationData[0].AuthorizationToken)
+	if err != nil {
+		return err
+	}
+	authInfoParts := strings.Split(string(authInfo), ":")
+	if len(authInfoParts) != 2 {
+		return errors.New("got malformed auth info from ECR")
+	}
+	dockerUsername := authInfoParts[0]
+	dockerPassword := authInfoParts[1]
+	sv.status("authenticating Docker client to ECR")
+	dockerLogin := exec.Command(
+		"docker", "login",
+		"--username", dockerUsername,
+		"--password-stdin",
+		fmt.Sprintf(
+			"%s.dkr.ecr.%s.amazonaws.com",
+			sv.awsAccountNumber, sv.awsRegion,
+		),
+	)
+	dockerLogin.Stdin = bytes.NewReader([]byte(dockerPassword))
+	dockerLogin.Stdout = os.Stdout
+	dockerLogin.Stderr = os.Stderr
+	if err := dockerLogin.Run(); err != nil {
+		return err
+	}
 	sv.status("downloading deployment config from S3")
 	dl := s3manager.NewDownloader(sv.s3)
 	buf := s3manager.NewWriteAtBuffer([]byte{})
@@ -424,6 +468,7 @@ func main() {
 		greenProxyHandler: httputil.NewSingleHostReverseProxy(greenUrl),
 		isGreen: isGreen,
 		s3: s3.NewFromConfig(awsCfg),
+		ecr: ecr.NewFromConfig(awsCfg),
 		awsRegion: awsCfg.Region,
 		awsAccountNumber: *ident.Account,
 		reloadJobs: map[string]*reloadJob{},
diff --git a/tf/iam.tf b/tf/iam.tf
index 88a5ec7..21d4a2a 100644
--- a/tf/iam.tf
+++ b/tf/iam.tf
@@ -49,6 +49,27 @@ data "aws_iam_policy_document" "server" {
       "arn:aws:s3:::${aws_s3_bucket.riju.bucket}/config.json",
     ]
   }
+
+  statement {
+    actions = [
+      "ecr:GetAuthorizationToken",
+    ]
+
+    resources = [
+      "*",
+    ]
+  }
+
+  statement {
+    actions = [
+      "ecr:BatchGetImage",
+      "ecr:GetDownloadUrlForLayer",
+    ]
+
+    resources = [
+      aws_ecr_repository.riju.arn,
+    ]
+  }
 }
 
 resource "aws_iam_policy" "server" {