From b2f85b70eab886bd3b3dd0410ecafb06261e4fd3 Mon Sep 17 00:00:00 2001
From: Vijay Krishnan <vijay.g.krishnan@oracle.com>
Date: Thu, 20 Sep 2018 13:57:41 -0700
Subject: [PATCH] Use registry auth token from Call extensions to pull images
 (#1228)

---
 api/agent/agent.go                 | 43 +++++++++++++++-------
 api/agent/agent_test.go            | 57 ++++++++++++++++++++++++++++++
 api/agent/drivers/docker/docker.go |  6 ++--
 3 files changed, 92 insertions(+), 14 deletions(-)

diff --git a/api/agent/agent.go b/api/agent/agent.go
index ba4962eb8..c26e5f6c7 100644
--- a/api/agent/agent.go
+++ b/api/agent/agent.go
@@ -21,6 +21,7 @@ import (
 	"github.com/fnproject/fn/api/models"
 	"github.com/fnproject/fn/fnext"
 	"github.com/fsnotify/fsnotify"
+	docker "github.com/fsouza/go-dockerclient"
 	"github.com/sirupsen/logrus"
 	"go.opencensus.io/stats"
 	"go.opencensus.io/trace"
@@ -118,6 +119,9 @@ type agent struct {
 // Option configures an agent at startup
 type Option func(*agent) error
 
+// RegistryToken is a reserved call extensions key to pass registry token
+const RegistryToken = "FN_REGISTRY_TOKEN"
+
 // New creates an Agent that executes functions locally as Docker containers.
 func New(da CallHandler, options ...Option) Agent {
 
@@ -869,14 +873,15 @@ func (a *agent) prepCold(ctx context.Context, call *call, tok ResourceToken, ch
 	}
 
 	container := &container{
-		id:      id.New().String(), // XXX we could just let docker generate ids...
-		image:   call.Image,
-		env:     map[string]string(call.Config),
-		memory:  call.Memory,
-		cpus:    uint64(call.CPUs),
-		fsSize:  a.cfg.MaxFsSize,
-		iofs:    &noopIOFS{},
-		timeout: time.Duration(call.Timeout) * time.Second, // this is unnecessary, but in case removal fails...
+		id:         id.New().String(), // XXX we could just let docker generate ids...
+		image:      call.Image,
+		env:        map[string]string(call.Config),
+		extensions: call.extensions,
+		memory:     call.Memory,
+		cpus:       uint64(call.CPUs),
+		fsSize:     a.cfg.MaxFsSize,
+		iofs:       &noopIOFS{},
+		timeout:    time.Duration(call.Timeout) * time.Second, // this is unnecessary, but in case removal fails...
 		logCfg: drivers.LoggerConfig{
 			URL: strings.TrimSpace(call.SyslogURL),
 			Tags: []drivers.LoggerTag{
@@ -1372,7 +1377,21 @@ func (c *container) WriteStat(ctx context.Context, stat drivers.Stat) {
 	c.swapMu.Unlock()
 }
 
-//func (c *container) DockerAuth() (docker.AuthConfiguration, error) {
-// Implementing the docker.AuthConfiguration interface.
-// TODO per call could implement this stored somewhere (vs. configured on host)
-//}
+// DockerAuth implements the docker.AuthConfiguration interface.
+func (c *container) DockerAuth() (*docker.AuthConfiguration, error) {
+	logger := common.Logger(context.TODO())
+	registryToken := ""
+	var ok bool
+	if registryToken, ok = c.extensions[RegistryToken]; !ok {
+		logger.WithField("Image", c.image).Infoln("No Registry Token for image")
+		registryToken = ""
+	} else {
+		logger.WithField("Image", c.image).Infof("Registry Token %s", registryToken)
+	}
+	if registryToken != "" {
+		return &docker.AuthConfiguration{
+			RegistryToken: registryToken,
+		}, nil
+	}
+	return nil, nil
+}
diff --git a/api/agent/agent_test.go b/api/agent/agent_test.go
index a2944e438..ecd24b4b1 100644
--- a/api/agent/agent_test.go
+++ b/api/agent/agent_test.go
@@ -1105,3 +1105,60 @@ func TestNBIOResourceTracker(t *testing.T) {
 		t.Fatalf("Expected successes, but got %d", ok)
 	}
 }
+
+func TestDockerAuthExtn(t *testing.T) {
+	modelCall := &models.Call{
+		AppID:       id.New().String(),
+		FnID:        id.New().String(),
+		Image:       "fnproject/fn-test-utils",
+		Type:        "sync",
+		Format:      "http",
+		Timeout:     1,
+		IdleTimeout: 2,
+	}
+	cfg, err := NewConfig()
+	if err != nil {
+		t.Fatalf("bad config %+v", cfg)
+	}
+
+	ls := logs.NewMock()
+	a := New(NewDirectCallDataAccess(ls, new(mqs.Mock)))
+	defer checkClose(t, a)
+
+	callIf, err := a.GetCall(FromModel(modelCall))
+	if err != nil {
+		t.Fatal(err)
+	}
+	call := callIf.(*call)
+
+	ctx := context.TODO()
+
+	c, err := newHotContainer(ctx, call, cfg)
+	if err != nil {
+		t.Fatal("got unexpected err: ", err)
+	}
+	da, err := c.DockerAuth()
+	if da != nil {
+		t.Fatal("invalid docker auth configuration")
+	}
+	if err != nil {
+		t.Fatal("got unexpected err: ", err)
+	}
+
+	// Add registry token as extension
+	extn := make(map[string]string)
+	extn["FN_REGISTRY_TOKEN"] = "TestRegistryToken"
+	call.extensions = extn
+
+	c, err = newHotContainer(ctx, call, cfg)
+	if err != nil {
+		t.Fatal("got unexpected err: ", err)
+	}
+	da, err = c.DockerAuth()
+	if da == nil {
+		t.Fatal("invalid docker auth configuration")
+	}
+	if da.RegistryToken != "TestRegistryToken" {
+		t.Fatalf("unexpected registry token %s", da.RegistryToken)
+	}
+}
diff --git a/api/agent/drivers/docker/docker.go b/api/agent/drivers/docker/docker.go
index 6542d11ad..93e421c03 100644
--- a/api/agent/drivers/docker/docker.go
+++ b/api/agent/drivers/docker/docker.go
@@ -36,7 +36,7 @@ type Auther interface {
 	// certain restrictions on images or if credentials must be acquired right
 	// before runtime and there's an error doing so. If these credentials don't
 	// work, the docker pull will fail and the task will be set to error status.
-	DockerAuth() (docker.AuthConfiguration, error)
+	DockerAuth() (*docker.AuthConfiguration, error)
 }
 
 type runResult struct {
@@ -310,7 +310,9 @@ func (drv *DockerDriver) ensureImage(ctx context.Context, task drivers.Container
 		if err != nil {
 			return err
 		}
-		config = &authConfig
+		if authConfig != nil {
+			config = authConfig
+		}
 	}
 
 	globalRepo := path.Join(reg, repo)