From 5a0f80391791770412ef96cd6ada21b2007dba9a Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Wed, 31 May 2017 15:36:04 -0700
Subject: [PATCH] imagesource: fill blob map on pull

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 snapshot/blobmapping/snapshotter.go | 21 +++++++++++++--------
 source/containerimage/pull.go       | 29 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/snapshot/blobmapping/snapshotter.go b/snapshot/blobmapping/snapshotter.go
index e0cdecbbf..b6bc377bc 100644
--- a/snapshot/blobmapping/snapshotter.go
+++ b/snapshot/blobmapping/snapshotter.go
@@ -53,6 +53,7 @@ func NewSnapshotter(opt Opt) (*Snapshotter, error) {
 	s := &Snapshotter{
 		Snapshotter: opt.Snapshotter,
 		db:          db,
+		opt:         opt,
 	}
 
 	return s, nil
@@ -67,7 +68,7 @@ func (s *Snapshotter) init() error {
 // Remove also removes a refrence to a blob. If it is a last reference then it deletes it the blob as well
 // Remove is not safe to be called concurrently
 func (s *Snapshotter) Remove(ctx context.Context, key string) error {
-	blob, err := s.Blob(ctx, key)
+	blob, err := s.GetBlob(ctx, key)
 	if err != nil {
 		return err
 	}
@@ -97,8 +98,8 @@ func (s *Snapshotter) Remove(ctx context.Context, key string) error {
 // TODO: make Blob/SetBlob part of generic metadata wrapper that can detect
 // blob key for deletion logic
 
-func (s *Snapshotter) Blob(ctx context.Context, key string) (string, error) {
-	var blob string
+func (s *Snapshotter) GetBlob(ctx context.Context, key string) (digest.Digest, error) {
+	var blob digest.Digest
 	err := s.db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket(bucketBySnapshot)
 		if b == nil {
@@ -106,7 +107,7 @@ func (s *Snapshotter) Blob(ctx context.Context, key string) (string, error) {
 		}
 		v := b.Get([]byte(key))
 		if v != nil {
-			blob = string(v)
+			blob = digest.Digest(v)
 		}
 		return nil
 	})
@@ -116,7 +117,11 @@ func (s *Snapshotter) Blob(ctx context.Context, key string) (string, error) {
 // Validates that there is no blob associated with the snapshot.
 // Checks that there is a blob in the content store.
 // If same blob has already been set then this is a noop.
-func (s *Snapshotter) SetBlob(ctx context.Context, key, blob string) error {
+func (s *Snapshotter) SetBlob(ctx context.Context, key string, blob digest.Digest) error {
+	_, err := s.opt.Content.Info(ctx, blob)
+	if err != nil {
+		return err
+	}
 	return s.db.Update(func(tx *bolt.Tx) error {
 		b, err := tx.CreateBucketIfNotExists(bucketBySnapshot)
 		if err != nil {
@@ -124,7 +129,7 @@ func (s *Snapshotter) SetBlob(ctx context.Context, key, blob string) error {
 		}
 		v := b.Get([]byte(key))
 		if v != nil {
-			if string(v) != blob {
+			if string(v) != string(blob) {
 				return errors.Errorf("different blob already set for %s", key)
 			} else {
 				return nil
@@ -142,8 +147,8 @@ func (s *Snapshotter) SetBlob(ctx context.Context, key, blob string) error {
 	})
 }
 
-func blobKey(blob, snapshot string) []byte {
-	return []byte(blob + "-" + snapshot)
+func blobKey(blob digest.Digest, snapshot string) []byte {
+	return []byte(string(blob) + "-" + snapshot)
 }
 
 // results are only valid for the lifetime of the transaction
diff --git a/source/containerimage/pull.go b/source/containerimage/pull.go
index 8a98a7061..b8de6cff5 100644
--- a/source/containerimage/pull.go
+++ b/source/containerimage/pull.go
@@ -11,6 +11,8 @@ import (
 	"github.com/containerd/containerd/remotes/docker"
 	"github.com/containerd/containerd/rootfs"
 	"github.com/containerd/containerd/snapshot"
+	digest "github.com/opencontainers/go-digest"
+	"github.com/opencontainers/image-spec/identity"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 	"github.com/tonistiigi/buildkit_poc/cache"
@@ -27,6 +29,11 @@ type SourceOpt struct {
 	CacheAccessor cache.Accessor
 }
 
+type blobmapper interface {
+	GetBlob(ctx context.Context, key string) (digest.Digest, error)
+	SetBlob(ctx context.Context, key string, blob digest.Digest) error
+}
+
 type imageSource struct {
 	SourceOpt
 	resolver remotes.Resolver
@@ -39,6 +46,11 @@ func NewSource(opt SourceOpt) (source.Source, error) {
 			Client: http.DefaultClient,
 		}),
 	}
+
+	if _, ok := opt.Snapshotter.(blobmapper); !ok {
+		return nil, errors.Errorf("imagesource requires snapshotter with blobs mapping support")
+	}
+
 	return is, nil
 }
 
@@ -93,9 +105,26 @@ func (is *imageSource) unpack(ctx context.Context, desc ocispec.Descriptor) (str
 	if err != nil {
 		return "", err
 	}
+
+	if err := is.fillBlobMapping(ctx, layers); err != nil {
+		return "", err
+	}
+
 	return string(chainID), nil
 }
 
+func (is *imageSource) fillBlobMapping(ctx context.Context, layers []rootfs.Layer) error {
+	var chain []digest.Digest
+	for _, l := range layers {
+		chain = append(chain, l.Diff.Digest)
+		chainID := identity.ChainID(chain)
+		if err := is.SourceOpt.Snapshotter.(blobmapper).SetBlob(ctx, string(chainID), l.Blob.Digest); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func getLayers(ctx context.Context, provider content.Provider, desc ocispec.Descriptor) ([]rootfs.Layer, error) {
 	p, err := content.ReadBlob(ctx, provider, desc.Digest)
 	if err != nil {