Added snapshot package and tests

executor/cmd/root.go

@@ -19,6 +19,7 @@ package cmd
 import (
 	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/constants"
 	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/dockerfile"
+	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/snapshot"
 	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/util"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -68,5 +69,13 @@ func execute() error {
 
 	// Unpack file system to root
 	logrus.Infof("Unpacking filesystem of %s...", baseImage)
-	return util.ExtractFileSystemFromImage(baseImage)
+	if err := util.ExtractFileSystemFromImage(baseImage); err != nil {
+		return err
+	}
+
+	l := snapshot.NewLayeredMap(util.Hasher())
+	snapshotter := snapshot.NewSnapshotter(l, constants.RootDir)
+
+	// Take initial snapshot
+	return snapshotter.Init()
 }

pkg/constants/constants.go

@@ -23,5 +23,8 @@ const (
 	// RootDir is the path to the root directory
 	RootDir = "/"
 
+	// WorkspaceDir is the path to the workspace directory
+	WorkspaceDir = "/workspace"
+
 	WhitelistPath = "/proc/self/mountinfo"
 )

pkg/snapshot/layered_map.go

@@ -0,0 +1,53 @@
+/*
+Copyright 2018 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package snapshot
+
+type LayeredMap struct {
+	layers []map[string]string
+	hasher func(string) string
+}
+
+func NewLayeredMap(h func(string) string) *LayeredMap {
+	l := LayeredMap{
+		hasher: h,
+	}
+	l.layers = []map[string]string{}
+	return &l
+}
+
+func (l *LayeredMap) Snapshot() {
+	l.layers = append(l.layers, map[string]string{})
+}
+
+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 {
+			return v, ok
+		}
+	}
+	return "", false
+}
+
+func (l *LayeredMap) MaybeAdd(s string) bool {
+	oldV, ok := l.Get(s)
+	newV := l.hasher(s)
+	if ok && newV == oldV {
+		return false
+	}
+	l.layers[len(l.layers)-1][s] = newV
+	return true
+}

pkg/snapshot/snapshot.go

@@ -0,0 +1,98 @@
+/*
+Copyright 2018 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package snapshot
+
+import (
+	"archive/tar"
+	"bytes"
+	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/util"
+	"github.com/sirupsen/logrus"
+
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+)
+
+// Snapshotter holds the root directory from which to take snapshots, and a list of snapshots taken
+type Snapshotter struct {
+	l         *LayeredMap
+	directory string
+	snapshots []string
+}
+
+// NewSnapshotter creates a new snapshotter rooted at d
+func NewSnapshotter(l *LayeredMap, d string) *Snapshotter {
+	return &Snapshotter{l: l, directory: d, snapshots: []string{}}
+}
+
+// Init initializes a new snapshotter
+func (s *Snapshotter) Init() error {
+	if _, err := s.snapShotFS(ioutil.Discard); err != nil {
+		return err
+	}
+	return nil
+}
+
+// TakeSnapshot takes a snapshot of the filesystem, avoiding directories in the whitelist
+// It stores changed files in a tar, and returns the contents of this tar at the end
+func (s *Snapshotter) TakeSnapshot() ([]byte, error) {
+
+	buf := bytes.NewBuffer([]byte{})
+	added, err := s.snapShotFS(buf)
+	if err != nil {
+		return nil, err
+	}
+	if !added {
+		logrus.Infof("No files were changed in this command, this layer will not be appended.")
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+	// Add buffer contents until buffer is empty
+	var contents []byte
+	for {
+		next := buf.Next(buf.Len())
+		if len(next) == 0 {
+			break
+		}
+		contents = append(contents, next...)
+	}
+	return contents, nil
+}
+
+func (s *Snapshotter) snapShotFS(f io.Writer) (bool, error) {
+	s.l.Snapshot()
+	added := false
+	w := tar.NewWriter(f)
+	defer w.Close()
+
+	err := filepath.Walk(s.directory, func(path string, info os.FileInfo, err error) error {
+		if util.PathInWhitelist(path, s.directory) {
+			return nil
+		}
+
+		// Only add to the tar if we add it to the layeredmap.
+		if s.l.MaybeAdd(path) {
+			added = true
+			return util.AddToTar(path, info, w)
+		}
+		return nil
+	})
+	return added, err
+}

pkg/snapshot/snapshot_test.go

@@ -0,0 +1,149 @@
+/*
+Copyright 2018 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package snapshot
+
+import (
+	"archive/tar"
+	"bytes"
+	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/util"
+	"github.com/GoogleCloudPlatform/k8s-container-builder/testutil"
+	"github.com/pkg/errors"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestSnapshotFileChange(t *testing.T) {
+
+	testDir, snapshotter, err := setUpTestDir()
+	defer os.RemoveAll(testDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Make some changes to the filesystem
+	newFiles := map[string]string{
+		"foo":           "newbaz1",
+		"workspace/bat": "bat",
+	}
+	if err := testutil.SetupFiles(testDir, newFiles); err != nil {
+		t.Fatalf("Error setting up fs: %s", err)
+	}
+	// Take another snapshot
+	contents, err := snapshotter.TakeSnapshot()
+	if err != nil {
+		t.Fatalf("Error taking snapshot of fs: %s", err)
+	}
+	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
+	reader := bytes.NewReader(contents)
+	tr := tar.NewReader(reader)
+	fooPath := filepath.Join(testDir, "foo")
+	snapshotFiles := map[string]string{
+		fooPath: "newbaz1",
+	}
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		if _, isFile := snapshotFiles[hdr.Name]; !isFile {
+			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
+		}
+		contents, _ := ioutil.ReadAll(tr)
+		if string(contents) != snapshotFiles[hdr.Name] {
+			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, snapshotFiles[hdr.Name], string(contents))
+		}
+	}
+}
+
+func TestSnapshotChangePermissions(t *testing.T) {
+	testDir, snapshotter, err := setUpTestDir()
+	defer os.RemoveAll(testDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Change permissions on a file
+	batPath := filepath.Join(testDir, "bar/bat")
+	if err := os.Chmod(batPath, 0600); err != nil {
+		t.Fatalf("Error changing permissions on %s: %v", batPath, err)
+	}
+	// Take another snapshot
+	contents, err := snapshotter.TakeSnapshot()
+	if err != nil {
+		t.Fatalf("Error taking snapshot of fs: %s", err)
+	}
+	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
+	reader := bytes.NewReader(contents)
+	tr := tar.NewReader(reader)
+	snapshotFiles := map[string]string{
+		batPath: "baz2",
+	}
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		if _, isFile := snapshotFiles[hdr.Name]; !isFile {
+			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
+		}
+		contents, _ := ioutil.ReadAll(tr)
+		if string(contents) != snapshotFiles[hdr.Name] {
+			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, snapshotFiles[hdr.Name], string(contents))
+		}
+	}
+}
+
+func TestEmptySnapshot(t *testing.T) {
+	testDir, snapshotter, err := setUpTestDir()
+	defer os.RemoveAll(testDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Take snapshot with no changes
+	contents, err := snapshotter.TakeSnapshot()
+	if err != nil {
+		t.Fatalf("Error taking snapshot of fs: %s", err)
+	}
+	// Since we took a snapshot with no changes, contents should be nil
+	if contents != nil {
+		t.Fatal("Contents should be nil, since no changes to the filesystem were made.")
+	}
+}
+
+func setUpTestDir() (string, *Snapshotter, error) {
+	testDir, err := ioutil.TempDir("", "")
+	if err != nil {
+		return testDir, nil, errors.Wrap(err, "setting up temp dir")
+	}
+	files := map[string]string{
+		"foo":            "baz1",
+		"bar/bat":        "baz2",
+		"workspace/file": "file",
+	}
+	// Set up initial files
+	if err := testutil.SetupFiles(testDir, files); err != nil {
+		return testDir, nil, errors.Wrap(err, "setting up file system")
+	}
+
+	// Take the initial snapshot
+	l := NewLayeredMap(util.Hasher())
+	snapshotter := NewSnapshotter(l, testDir)
+	if err := snapshotter.Init(); err != nil {
+		return testDir, nil, errors.Wrap(err, "initializing snapshotter")
+	}
+	return testDir, snapshotter, nil
+}

pkg/util/fs_util.go

@@ -24,6 +24,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"io"
 	"os"
+	"path/filepath"
 	"strings"
 )
 
@@ -47,6 +48,17 @@ func ExtractFileSystemFromImage(img string) error {
 	return pkgutil.GetFileSystemFromReference(ref, imgSrc, constants.RootDir, whitelist)
 }
 
+// PathInWhitelist returns true if the path is whitelisted
+func PathInWhitelist(path, directory string) bool {
+	for _, d := range whitelist {
+		dirPath := filepath.Join(directory, d)
+		if pkgutil.HasFilepathPrefix(path, dirPath) {
+			return true
+		}
+	}
+	return false
+}
+
 // Get whitelist from roots of mounted files
 // Each line of /proc/self/mountinfo is in the form:
 // 36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue