Add whiteout handling by switching to a two-phase approach. (#139)

* Add whiteout handling by switching to a two-phase approach.

Also only handle hardlinks within one layer

* Simplify the run test.

pkg/snapshot/layered_map.go

@@ -16,6 +16,11 @@ limitations under the License.
 
 package snapshot
 
+import (
+	"path/filepath"
+	"strings"
+)
+
 type LayeredMap struct {
 	layers []map[string]string
 	hasher func(string) (string, error)
@@ -33,6 +38,21 @@ func (l *LayeredMap) Snapshot() {
 	l.layers = append(l.layers, map[string]string{})
 }
 
+func (l *LayeredMap) GetFlattenedPathsForWhiteOut() map[string]struct{} {
+	paths := map[string]struct{}{}
+	for _, l := range l.layers {
+		for p := range l {
+			if strings.HasPrefix(filepath.Base(p), ".wh.") {
+				delete(paths, p)
+			} else {
+				paths[p] = struct{}{}
+			}
+			paths[p] = struct{}{}
+		}
+	}
+	return paths
+}
+
 func (l *LayeredMap) Get(s string) (string, bool) {
 	for i := len(l.layers) - 1; i >= 0; i-- {
 		if v, ok := l.layers[i][s]; ok {

pkg/snapshot/snapshot.go

@@ -33,6 +33,7 @@ import (
 type Snapshotter struct {
 	l         *LayeredMap
 	directory string
+	hardlinks map[uint64]string
 }
 
 // NewSnapshotter creates a new snapshotter rooted at d
@@ -60,10 +61,7 @@ func (s *Snapshotter) TakeSnapshot(files []string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	contents, err := ioutil.ReadAll(buf)
-	if err != nil {
-		return nil, err
-	}
+	contents := buf.Bytes()
 	if !filesAdded {
 		return nil, nil
 	}
@@ -99,7 +97,7 @@ func (s *Snapshotter) TakeSnapshotOfFiles(files []string) ([]byte, error) {
 		}
 		if maybeAdd {
 			filesAdded = true
-			util.AddToTar(file, info, w)
+			util.AddToTar(file, info, s.hardlinks, w)
 		}
 	}
 	if !filesAdded {
@@ -109,27 +107,55 @@ func (s *Snapshotter) TakeSnapshotOfFiles(files []string) ([]byte, error) {
 }
 
 func (s *Snapshotter) snapShotFS(f io.Writer) (bool, error) {
+	s.hardlinks = map[uint64]string{}
 	s.l.Snapshot()
+	existingPaths := s.l.GetFlattenedPathsForWhiteOut()
 	filesAdded := false
 	w := tar.NewWriter(f)
 	defer w.Close()
 
-	err := filepath.Walk(s.directory, func(path string, info os.FileInfo, err error) error {
+	// Save the fs state in a map to iterate over later.
+	memFs := map[string]os.FileInfo{}
+	filepath.Walk(s.directory, func(path string, info os.FileInfo, err error) error {
+		memFs[path] = info
+		return nil
+	})
+
+	// First handle whiteouts
+	for p := range memFs {
+		delete(existingPaths, p)
+	}
+	for path := range existingPaths {
+		// Only add the whiteout if the directory for the file still exists.
+		dir := filepath.Dir(path)
+		if _, ok := memFs[dir]; ok {
+			logrus.Infof("Adding whiteout for %s", path)
+			filesAdded = true
+			if err := util.Whiteout(path, w); err != nil {
+				return false, err
+			}
+		}
+	}
+
+	// Now create the tar.
+	for path, info := range memFs {
 		if util.PathInWhitelist(path, s.directory) {
 			logrus.Debugf("Not adding %s to layer, as it's whitelisted", path)
-			return nil
+			continue
 		}
 
 		// Only add to the tar if we add it to the layeredmap.
 		maybeAdd, err := s.l.MaybeAdd(path)
 		if err != nil {
-			return err
+			return false, err
 		}
 		if maybeAdd {
 			filesAdded = true
-			return util.AddToTar(path, info, w)
+			if err := util.AddToTar(path, info, s.hardlinks, w); err != nil {
+				return false, err
+			}
 		}
-		return nil
-	})
-	return filesAdded, err
+	}
+
+	return filesAdded, nil
 }

pkg/util/tar_util.go

@@ -23,6 +23,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 	"syscall"
 
 	pkgutil "github.com/GoogleContainerTools/container-diff/pkg/util"
@@ -31,10 +32,8 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-var hardlinks = make(map[uint64]string)
-
 // AddToTar adds the file i to tar w at path p
-func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
+func AddToTar(p string, i os.FileInfo, hardlinks map[uint64]string, w *tar.Writer) error {
 	linkDst := ""
 	if i.Mode()&os.ModeSymlink != 0 {
 		var err error
@@ -49,7 +48,7 @@ func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
 	}
 	hdr.Name = p
 
-	hardlink, linkDst := checkHardlink(p, i)
+	hardlink, linkDst := checkHardlink(p, hardlinks, i)
 	if hardlink {
 		hdr.Linkname = linkDst
 		hdr.Typeflag = tar.TypeLink
@@ -72,8 +71,23 @@ func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
 	return nil
 }
 
+func Whiteout(p string, w *tar.Writer) error {
+	dir := filepath.Dir(p)
+	name := ".wh." + filepath.Base(p)
+
+	th := &tar.Header{
+		Name: filepath.Join(dir, name),
+		Size: 0,
+	}
+	if err := w.WriteHeader(th); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Returns true if path is hardlink, and the link destination
-func checkHardlink(p string, i os.FileInfo) (bool, string) {
+func checkHardlink(p string, hardlinks map[uint64]string, i os.FileInfo) (bool, string) {
 	hardlink := false
 	linkDst := ""
 	if sys := i.Sys(); sys != nil {

pkg/util/tar_util_test.go

@@ -101,7 +101,7 @@ func createTar(testdir string, writer io.Writer) error {
 		if err != nil {
 			return err
 		}
-		if err := AddToTar(filePath, fi, w); err != nil {
+		if err := AddToTar(filePath, fi, map[uint64]string{}, w); err != nil {
 			return err
 		}
 	}