From 73e37f9596280520020af5eb093b4442f9950c51 Mon Sep 17 00:00:00 2001
From: cbkhare <Cbkhare@users.noreply.github.com>
Date: Mon, 8 Mar 2021 18:27:43 +0530
Subject: [PATCH] docker graceful exits when image pull fails #382

added check for Decode statement

added unit tests

added unit tests

added unit tests

added unit tests
---
 cmd/tink-worker/internal/registry.go      | 32 ++++++++--
 cmd/tink-worker/internal/registry_test.go | 75 +++++++++++++++++++++++
 2 files changed, 103 insertions(+), 4 deletions(-)
 create mode 100644 cmd/tink-worker/internal/registry_test.go

diff --git a/cmd/tink-worker/internal/registry.go b/cmd/tink-worker/internal/registry.go
index a20f12e32..97bc13ad5 100644
--- a/cmd/tink-worker/internal/registry.go
+++ b/cmd/tink-worker/internal/registry.go
@@ -5,7 +5,6 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"io"
-	"os"
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
@@ -22,6 +21,17 @@ type RegistryConnDetails struct {
 	logger log.Logger
 }
 
+// ImagePullStatus is the status of the downloaded Image chunk
+type ImagePullStatus struct {
+	Status         string `json:"status"`
+	Error          string `json:"error"`
+	Progress       string `json:"progress"`
+	ProgressDetail struct {
+		Current int `json:"current"`
+		Total   int `json:"total"`
+	} `json:"progressDetail"`
+}
+
 // NewRegistryConnDetails creates a new RegistryConnDetails
 func NewRegistryConnDetails(registry, user, pwd string, logger log.Logger) *RegistryConnDetails {
 	return &RegistryConnDetails{
@@ -46,8 +56,12 @@ func (r *RegistryConnDetails) NewClient() (*client.Client, error) {
 	return c, nil
 }
 
+type imagePuller interface {
+	ImagePull(context.Context, string, types.ImagePullOptions) (io.ReadCloser, error)
+}
+
 // pullImage outputs to stdout the contents of the requested image (relative to the registry)
-func (r *RegistryConnDetails) pullImage(ctx context.Context, cli *client.Client, image string) error {
+func (r *RegistryConnDetails) pullImage(ctx context.Context, cli imagePuller, image string) error {
 	authConfig := types.AuthConfig{
 		Username:      r.user,
 		Password:      r.pwd,
@@ -64,8 +78,18 @@ func (r *RegistryConnDetails) pullImage(ctx context.Context, cli *client.Client,
 		return errors.Wrap(err, "DOCKER PULL")
 	}
 	defer out.Close()
-	if _, err := io.Copy(os.Stdout, out); err != nil {
-		return err
+	fd := json.NewDecoder(out)
+	var status *ImagePullStatus
+	for {
+		if err := fd.Decode(&status); err != nil {
+			if err == io.EOF {
+				break
+			}
+			return errors.Wrap(err, "DOCKER PULL")
+		}
+		if status.Error != "" {
+			return errors.Wrap(errors.New(status.Error), "DOCKER PULL")
+		}
 	}
 	return nil
 }
diff --git a/cmd/tink-worker/internal/registry_test.go b/cmd/tink-worker/internal/registry_test.go
new file mode 100644
index 000000000..b6ab5c783
--- /dev/null
+++ b/cmd/tink-worker/internal/registry_test.go
@@ -0,0 +1,75 @@
+package internal
+
+import (
+	"context"
+	"errors"
+	"io"
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"github.com/docker/docker/api/types"
+	"github.com/packethost/pkg/log"
+	//"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func setupTestLogger(t *testing.T) log.Logger {
+	service := "github.com/tinkerbell/tink"
+	logger, err := log.Init(service)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return logger
+}
+
+type imagePullerMock struct {
+	stringReadCloser io.ReadCloser
+	imagePullErr     error
+}
+
+func (d *imagePullerMock) ImagePull(ctx context.Context, str string, op types.ImagePullOptions) (io.ReadCloser, error) {
+	return d.stringReadCloser, d.imagePullErr
+}
+
+func TestPullImageAnyFailure(t *testing.T) {
+	for _, test := range []struct {
+		testName         string
+		testString       string
+		testImagePullErr error
+		testErr          error
+	}{
+		{
+			testName:         "success",
+			testString:       "{\"status\": \"hello\",\"error\":\"\"}{\"status\":\"world\",\"error\":\"\"}",
+			testImagePullErr: nil,
+			testErr:          nil,
+		},
+		{
+			testName:         "fail",
+			testString:       "{\"error\": \"\"}",
+			testImagePullErr: errors.New("Tested, failure of the image pull"),
+			testErr:          errors.New("DOCKER PULL: Tested, failure of the image pull"),
+		},
+		{
+			testName:         "fail_partial",
+			testString:       "{\"status\": \"hello\",\"error\":\"\"}{\"status\":\"world\",\"error\":\"Tested, failure of No space left on device\"}",
+			testImagePullErr: nil,
+			testErr:          errors.New("DOCKER PULL: Tested, failure of No space left on device"),
+		},
+	} {
+		t.Run(test.testName, func(t *testing.T) {
+			ctx := context.Background()
+			rcon := NewRegistryConnDetails("test", "testUser", "testPwd", setupTestLogger(t))
+			stringReader := strings.NewReader(test.testString)
+			cli := &imagePullerMock{
+				stringReadCloser: ioutil.NopCloser(stringReader),
+				imagePullErr:     test.testImagePullErr,
+			}
+			err := rcon.pullImage(ctx, cli, test.testName)
+			if test.testErr != nil {
+				assert.Equal(t, err.Error(), test.testErr.Error())
+			}
+		})
+	}
+}