From 13d585059f1a10233b9c985d47d637cc9b129a8c Mon Sep 17 00:00:00 2001
From: Brad Davidson <brad.davidson@rancher.com>
Date: Mon, 14 Dec 2020 16:26:14 -0800
Subject: [PATCH] Add registry mirrors for CI test step

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 .drone.yml             |  3 +++
 Dockerfile.test.dapper |  4 ++--
 scripts/test-helpers   | 50 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/.drone.yml b/.drone.yml
index 28b4b5a7a6..57fc5e6f1b 100644
--- a/.drone.yml
+++ b/.drone.yml
@@ -65,6 +65,7 @@ steps:
   image: rancher/dapper:v0.5.0
   secrets: [ gcloud_auth ]
   environment:
+    ENABLE_REGISTRY: 'true'
     GCLOUD_AUTH:
       from_secret: gcloud_auth
   commands:
@@ -160,6 +161,7 @@ steps:
   image: rancher/dapper:v0.5.0
   secrets: [ gcloud_auth ]
   environment:
+    ENABLE_REGISTRY: 'true'
     GCLOUD_AUTH:
       from_secret: gcloud_auth
   commands:
@@ -238,6 +240,7 @@ steps:
   image: rancher/dapper:v0.5.0
   secrets: [ gcloud_auth ]
   environment:
+    ENABLE_REGISTRY: 'true'
     GCLOUD_AUTH:
       from_secret: gcloud_auth
   commands:
diff --git a/Dockerfile.test.dapper b/Dockerfile.test.dapper
index 2eb4d98747..ff9d340e70 100644
--- a/Dockerfile.test.dapper
+++ b/Dockerfile.test.dapper
@@ -16,8 +16,8 @@ RUN OS=linux; \
 
 ENV TEST_CLEANUP true
 
-ENV DAPPER_RUN_ARGS --privileged --network host
-ENV DAPPER_ENV REPO TAG DRONE_TAG DRONE_BUILD_EVENT IMAGE_NAME GCLOUD_AUTH SONOBUOY_VERSION
+ENV DAPPER_RUN_ARGS --privileged --network host -v /tmp:/tmp
+ENV DAPPER_ENV REPO TAG DRONE_TAG DRONE_BUILD_EVENT IMAGE_NAME GCLOUD_AUTH SONOBUOY_VERSION ENABLE_REGISTRY
 ENV DAPPER_SOURCE /go/src/github.com/rancher/k3s/
 ENV DAPPER_OUTPUT ./dist
 ENV DAPPER_DOCKER_SOCKET true
diff --git a/scripts/test-helpers b/scripts/test-helpers
index 449bb070cf..3b626b607c 100755
--- a/scripts/test-helpers
+++ b/scripts/test-helpers
@@ -419,6 +419,7 @@ provision-server() {
         -p 6443 \
         -e K3S_TOKEN=$(cat $TEST_DIR/metadata/secret) \
         -e K3S_DEBUG=true \
+        ${REGISTRY_CLUSTER_ARGS:-} \
         $K3S_IMAGE server $ARGS $SERVER_ARGS ${!SERVER_INSTANCE_ARGS}
 
     local ip=$(docker inspect --format '{{ .NetworkSettings.IPAddress }}' $name | tee $TEST_DIR/servers/$count/metadata/ip)
@@ -445,6 +446,7 @@ provision-agent() {
         --privileged \
         -e K3S_TOKEN=$(cat $TEST_DIR/metadata/secret) \
         -e K3S_URL=$K3S_URL \
+        ${REGISTRY_CLUSTER_ARGS:-} \
         $K3S_IMAGE agent $ARGS $AGENT_ARGS ${!AGENT_INSTANCE_ARGS}
 
     echo "Started $name"
@@ -457,6 +459,10 @@ export -f provision-agent
 provision-cluster() {
     run-function cluster-pre-hook
 
+    if [ "${ENABLE_REGISTRY}" == 'true' ]; then
+      provision-registry-proxy
+    fi
+
     for i in $(seq 1 $NUM_SERVERS); do
         provision-server
         timeout --foreground 120s bash -c "wait-for-kubeconfig $i"
@@ -480,6 +486,50 @@ export -f provision-cluster
 
 # ---
 
+provision-registry-proxy() {
+    set -e -o pipefail
+    local image="docker.io/library/registry:2.7.1"
+    local prefix="docker-registry-"
+    local registries="docker.io:registry-1.docker.io k8s.gcr.io gcr.io quay.io ghcr.io"
+    local registries_yaml="$TEST_DIR/registries.yaml"
+
+    echo "mirrors:" > $registries_yaml
+
+    for registry in $registries; do
+        IFS=: read registry_name registry_endpoint <<< $registry
+        if [ -z "$registry_endpoint" ]; then
+            registry_endpoint=$registry_name
+        fi
+
+        local name="registry_${registry_name//./_}"
+        local status=$(docker inspect $name --format '{{ .State.Status }} {{ .Config.Image }} {{ (index .HostConfig.PortBindings "5000/tcp" 0).HostPort }}' 2>/dev/null || true)
+
+        read state_status config_image hostport <<< $status
+        if [ "$state_status" != "running" ] || [ "$config_image" != "$image" ]; then
+            hostport=$(timeout --foreground 5s bash -c get-port)
+            docker rm --force $name 2>/dev/null || true
+            docker run \
+                -d --name $name \
+                -p 0.0.0.0:$hostport:5000 \
+                -v "registry-cache:/var/lib/registry" \
+                -e "REGISTRY_HTTP_SECRET=shared-secret" \
+                -e "REGISTRY_PROXY_REMOTEURL=https://$registry_endpoint" \
+                -e "REGISTRY_STORAGE_CACHE_BLOBDESCRIPTOR=inmemory" \
+                -e "REGISTRY_STORAGE_FILESYSTEM_ROOTDIRECTORY=/var/lib/registry/$registry_name" \
+                $image
+        fi
+        echo -e "  $registry_name:\n    endpoint:\n    - http://172.17.0.1:$hostport" >> $registries_yaml
+    done
+
+    echo "Using registry mirror with cluster registries.yaml:"
+    cat $registries_yaml
+
+    export REGISTRY_CLUSTER_ARGS="-v $registries_yaml:/etc/rancher/k3s/registries.yaml"
+}
+export -f provision-registry-proxy
+
+# ---
+
 early-exit() {
     printf "\033[33m$1\033[m\n"
     exit $2