fix 8926: rmi dangling is unsafe when pulling

Signed-off-by: Ma Shimiao <mashimiao.fnst@cn.fujitsu.com>
Signed-off-by: Tibor Vass <tibor@docker.com>

builder/evaluator.go

@@ -131,6 +131,9 @@ type Builder struct {
 	memorySwap   int64
 
 	cancelled <-chan struct{} // When closed, job was cancelled.
+
+	activeImages []string
+	id           string // Used to hold reference images
 }
 
 // Run the builder with the context. This is the lynchpin of this package. This

builder/internals.go

@@ -128,6 +128,8 @@ func (b *Builder) commit(id string, autoCmd *runconfig.Command, comment string)
 	if err != nil {
 		return err
 	}
+	b.Daemon.Graph().Retain(b.id, image.ID)
+	b.activeImages = append(b.activeImages, image.ID)
 	b.image = image.ID
 	return nil
 }
@@ -592,6 +594,8 @@ func (b *Builder) probeCache() (bool, error) {
 	fmt.Fprintf(b.OutStream, " ---> Using cache\n")
 	logrus.Debugf("[BUILDER] Use cached version")
 	b.image = cache.ID
+	b.Daemon.Graph().Retain(b.id, cache.ID)
+	b.activeImages = append(b.activeImages, cache.ID)
 	return true, nil
 }
 

builder/job.go

@@ -20,6 +20,7 @@ import (
 	"github.com/docker/docker/pkg/parsers"
 	"github.com/docker/docker/pkg/progressreader"
 	"github.com/docker/docker/pkg/streamformatter"
+	"github.com/docker/docker/pkg/stringid"
 	"github.com/docker/docker/pkg/urlutil"
 	"github.com/docker/docker/registry"
 	"github.com/docker/docker/runconfig"
@@ -200,13 +201,17 @@ func Build(d *daemon.Daemon, buildConfig *Config) error {
 		memory:          buildConfig.Memory,
 		memorySwap:      buildConfig.MemorySwap,
 		cancelled:       buildConfig.WaitCancelled(),
+		id:              stringid.GenerateRandomID(),
 	}
 
+	defer func() {
+		builder.Daemon.Graph().Release(builder.id, builder.activeImages...)
+	}()
+
 	id, err := builder.Run(context)
 	if err != nil {
 		return err
 	}
-
 	if repoName != "" {
 		return d.Repositories().Tag(repoName, tag, id, true)
 	}

daemon/image_delete.go

@@ -137,6 +137,9 @@ func (daemon *Daemon) imgDeleteHelper(name string, list *[]types.ImageDelete, fi
 }
 
 func (daemon *Daemon) canDeleteImage(imgID string, force bool) error {
+	if daemon.Graph().IsHeld(imgID) {
+		return fmt.Errorf("Conflict, cannot delete because %s is held by an ongoing pull or build", stringid.TruncateID(imgID))
+	}
 	for _, container := range daemon.List() {
 		if container.ImageID == "" {
 			// This technically should never happen, but if the container

graph/graph.go

@@ -13,6 +13,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/Sirupsen/logrus"
@@ -29,12 +30,56 @@ import (
 	"github.com/docker/docker/runconfig"
 )
 
+// The type is used to protect pulling or building related image
+// layers from deleteing when filtered by dangling=true
+// The key of layers is the images ID which is pulling or building
+// The value of layers is a slice which hold layer IDs referenced to
+// pulling or building images
+type retainedLayers struct {
+	layerHolders map[string]map[string]struct{} // map[layerID]map[sessionID]
+	sync.Mutex
+}
+
+func (r *retainedLayers) Add(sessionID string, layerIDs []string) {
+	r.Lock()
+	defer r.Unlock()
+	for _, layerID := range layerIDs {
+		if r.layerHolders[layerID] == nil {
+			r.layerHolders[layerID] = map[string]struct{}{}
+		}
+		r.layerHolders[layerID][sessionID] = struct{}{}
+	}
+}
+
+func (r *retainedLayers) Delete(sessionID string, layerIDs []string) {
+	r.Lock()
+	defer r.Unlock()
+	for _, layerID := range layerIDs {
+		holders, ok := r.layerHolders[layerID]
+		if !ok {
+			continue
+		}
+		delete(holders, sessionID)
+		if len(holders) == 0 {
+			delete(r.layerHolders, layerID) // Delete any empty reference set.
+		}
+	}
+}
+
+func (r *retainedLayers) Exists(layerID string) bool {
+	r.Lock()
+	_, exists := r.layerHolders[layerID]
+	r.Unlock()
+	return exists
+}
+
 // A Graph is a store for versioned filesystem images and the relationship between them.
 type Graph struct {
 	root       string
 	idIndex    *truncindex.TruncIndex
 	driver     graphdriver.Driver
 	imageMutex imageMutex // protect images in driver.
+	retained   *retainedLayers
 }
 
 type Image struct {
@@ -76,6 +121,7 @@ func NewGraph(root string, driver graphdriver.Driver) (*Graph, error) {
 		root:     abspath,
 		idIndex:  truncindex.NewTruncIndex([]string{}),
 		driver:   driver,
+		retained: &retainedLayers{layerHolders: make(map[string]map[string]struct{})},
 	}
 	if err := graph.restore(); err != nil {
 		return nil, err
@@ -83,6 +129,11 @@ func NewGraph(root string, driver graphdriver.Driver) (*Graph, error) {
 	return graph, nil
 }
 
+// IsHeld returns whether the given layerID is being used by an ongoing pull or build.
+func (graph *Graph) IsHeld(layerID string) bool {
+	return graph.retained.Exists(layerID)
+}
+
 func (graph *Graph) restore() error {
 	dir, err := ioutil.ReadDir(graph.root)
 	if err != nil {
@@ -367,6 +418,17 @@ func (graph *Graph) ByParent() map[string][]*Image {
 	return byParent
 }
 
+// If the images and layers are in pulling chain, retain them.
+// If not, they may be deleted by rmi with dangling condition.
+func (graph *Graph) Retain(sessionID string, layerIDs ...string) {
+	graph.retained.Add(sessionID, layerIDs)
+}
+
+// Release removes the referenced image id from the provided set of layers.
+func (graph *Graph) Release(sessionID string, layerIDs ...string) {
+	graph.retained.Delete(sessionID, layerIDs)
+}
+
 // Heads returns all heads in the graph, keyed by id.
 // A head is an image which is not the parent of another image in the graph.
 func (graph *Graph) Heads() map[string]*Image {

graph/pull_v1.go

@@ -117,6 +117,11 @@ func (p *v1Puller) pullRepository(askedTag string) error {
 	errors := make(chan error)
 
 	layersDownloaded := false
+	imgIDs := []string{}
+	sessionID := p.session.ID()
+	defer func() {
+		p.graph.Release(sessionID, imgIDs...)
+	}()
 	for _, image := range repoData.ImgList {
 		downloadImage := func(img *registry.ImgData) {
 			if askedTag != "" && img.Tag != askedTag {
@@ -144,6 +149,10 @@ func (p *v1Puller) pullRepository(askedTag string) error {
 			}
 			defer p.poolRemove("pull", "img:"+img.ID)
 
+			// we need to retain it until tagging
+			p.graph.Retain(sessionID, img.ID)
+			imgIDs = append(imgIDs, img.ID)
+
 			out.Write(p.sf.FormatProgress(stringid.TruncateID(img.ID), fmt.Sprintf("Pulling image (%s) from %s", img.Tag, p.repoInfo.CanonicalName), nil))
 			success := false
 			var lastErr, err error
@@ -226,6 +235,11 @@ func (p *v1Puller) pullImage(imgID, endpoint string, token []string) (bool, erro
 	// FIXME: Try to stream the images?
 	// FIXME: Launch the getRemoteImage() in goroutines
 
+	sessionID := p.session.ID()
+	// As imgID has been retained in pullRepository, no need to retain again
+	p.graph.Retain(sessionID, history[1:]...)
+	defer p.graph.Release(sessionID, history[1:]...)
+
 	layersDownloaded := false
 	for i := len(history) - 1; i >= 0; i-- {
 		id := history[i]

graph/pull_v2.go

@@ -26,6 +26,7 @@ type v2Puller struct {
 	sf        *streamformatter.StreamFormatter
 	repoInfo  *registry.RepositoryInfo
 	repo      distribution.Repository
+	sessionID string
 }
 
 func (p *v2Puller) Pull(tag string) (fallback bool, err error) {
@@ -36,6 +37,8 @@ func (p *v2Puller) Pull(tag string) (fallback bool, err error) {
 		return true, err
 	}
 
+	p.sessionID = stringid.GenerateRandomID()
+
 	if err := p.pullV2Repository(tag); err != nil {
 		if registry.ContinueOnError(err) {
 			logrus.Debugf("Error trying v2 registry: %v", err)
@@ -198,6 +201,12 @@ func (p *v2Puller) pullV2Tag(tag, taggedName string) (bool, error) {
 	out.Write(p.sf.FormatStatus(tag, "Pulling from %s", p.repo.Name()))
 
 	downloads := make([]downloadInfo, len(manifest.FSLayers))
+
+	layerIDs := []string{}
+	defer func() {
+		p.graph.Release(p.sessionID, layerIDs...)
+	}()
+
 	for i := len(manifest.FSLayers) - 1; i >= 0; i-- {
 		img, err := NewImgJSON([]byte(manifest.History[i].V1Compatibility))
 		if err != nil {
@@ -207,6 +216,9 @@ func (p *v2Puller) pullV2Tag(tag, taggedName string) (bool, error) {
 		downloads[i].img = img
 		downloads[i].digest = manifest.FSLayers[i].BlobSum
 
+		p.graph.Retain(p.sessionID, img.ID)
+		layerIDs = append(layerIDs, img.ID)
+
 		// Check if exists
 		if p.graph.Exists(img.ID) {
 			logrus.Debugf("Image already exists: %s", img.ID)

registry/session.go

@@ -23,6 +23,7 @@ import (
 	"github.com/docker/docker/cliconfig"
 	"github.com/docker/docker/pkg/httputils"
 	"github.com/docker/docker/pkg/ioutils"
+	"github.com/docker/docker/pkg/stringid"
 	"github.com/docker/docker/pkg/tarsum"
 )
 
@@ -35,6 +36,7 @@ type Session struct {
 	client        *http.Client
 	// TODO(tiborvass): remove authConfig
 	authConfig *cliconfig.AuthConfig
+	id         string
 }
 
 type authTransport struct {
@@ -158,6 +160,7 @@ func NewSession(client *http.Client, authConfig *cliconfig.AuthConfig, endpoint
 		authConfig:    authConfig,
 		client:        client,
 		indexEndpoint: endpoint,
+		id:            stringid.GenerateRandomID(),
 	}
 
 	var alwaysSetBasicAuth bool
@@ -188,6 +191,11 @@ func NewSession(client *http.Client, authConfig *cliconfig.AuthConfig, endpoint
 	return r, nil
 }
 
+// ID returns this registry session's ID.
+func (r *Session) ID() string {
+	return r.id
+}
+
 // Retrieve the history of a given image from the Registry.
 // Return a list of the parent's json (requested image included)
 func (r *Session) GetRemoteHistory(imgID, registry string) ([]string, error) {