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

pkg/util/tar_util.go

@@ -0,0 +1,52 @@
+/*
+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 util
+
+import (
+	"archive/tar"
+	"io"
+	"os"
+)
+
+// AddToTar adds the file i to tar w at path p
+func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
+	linkDst := ""
+	if i.Mode()&os.ModeSymlink != 0 {
+		var err error
+		linkDst, err = os.Readlink(p)
+		if err != nil {
+			return err
+		}
+	}
+	hdr, err := tar.FileInfoHeader(i, linkDst)
+	if err != nil {
+		return err
+	}
+	hdr.Name = p
+	w.WriteHeader(hdr)
+	if !i.Mode().IsRegular() {
+		return nil
+	}
+	r, err := os.Open(p)
+	if err != nil {
+		return err
+	}
+	if _, err := io.Copy(w, r); err != nil {
+		return err
+	}
+	return nil
+}

pkg/util/util.go

@@ -17,8 +17,12 @@ limitations under the License.
 package util
 
 import (
+	"crypto/md5"
+	"encoding/hex"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
+	"io"
+	"os"
 )
 
 // SetLogLevel sets the logrus logging level
@@ -30,3 +34,30 @@ func SetLogLevel(logLevel string) error {
 	logrus.SetLevel(lvl)
 	return nil
 }
+
+// Hasher returns a hash function, used in snapshotting to determine if a file has changed
+func Hasher() func(string) string {
+	hasher := func(p string) string {
+		h := md5.New()
+		fi, err := os.Lstat(p)
+		if err != nil {
+			panic(err)
+		}
+		h.Write([]byte(fi.Mode().String()))
+		h.Write([]byte(fi.ModTime().String()))
+
+		if fi.Mode().IsRegular() {
+			f, err := os.Open(p)
+			if err != nil {
+				panic(err)
+			}
+			defer f.Close()
+			if _, err := io.Copy(h, f); err != nil {
+				panic(err)
+			}
+		}
+
+		return hex.EncodeToString(h.Sum(nil))
+	}
+	return hasher
+}

testutil/util.go

@@ -18,10 +18,27 @@ package testutil
 
 import (
 	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 	"reflect"
 	"testing"
 )
 
+// SetupFiles creates files at path
+func SetupFiles(path string, files map[string]string) error {
+	for p, c := range files {
+		path := filepath.Join(path, p)
+		if err := os.MkdirAll(filepath.Dir(path), 0750); err != nil {
+			return err
+		}
+		if err := ioutil.WriteFile(path, []byte(c), 0644); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func CheckErrorAndDeepEqual(t *testing.T, shouldErr bool, err error, expected, actual interface{}) {
 	if err := checkErr(shouldErr, err); err != nil {
 		t.Error(err)