Merge pull request #27 from priyawadhwa/snapshot

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()
 }

integration_tests/dockerfiles/config_test_extract_fs.json

@@ -7,7 +7,7 @@
       "Adds": [
         {
           "Name": "/workspace",
-          "Size": 16271436
+          "Size": 16298144
         },
         {
           "Name": "/workspace/Dockerfile",
@@ -15,7 +15,7 @@
         },
         {
           "Name": "/workspace/executor",
-          "Size": 16271436
+          "Size": 16298144
         }
       ],
       "Dels": [

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,56 @@
+/*
+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, error)
+}
+
+func NewLayeredMap(h func(string) (string, error)) *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, error) {
+	oldV, ok := l.Get(s)
+	newV, err := l.hasher(s)
+	if err != nil {
+		return false, err
+	}
+	if ok && newV == oldV {
+		return false, nil
+	}
+	l.layers[len(l.layers)-1][s] = newV
+	return true, nil
+}

pkg/snapshot/snapshot.go

@@ -0,0 +1,122 @@
+/*
+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
+}
+
+// NewSnapshotter creates a new snapshotter rooted at d
+func NewSnapshotter(l *LayeredMap, d string) *Snapshotter {
+	return &Snapshotter{l: l, directory: d}
+}
+
+// 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, and creates
+// a tarball of the changed files. Return contents of the tarball, and whether or not any files were changed
+func (s *Snapshotter) TakeSnapshot() ([]byte, bool, error) {
+	buf := bytes.NewBuffer([]byte{})
+	filesAdded, err := s.snapShotFS(buf)
+	if err != nil {
+		return nil, filesAdded, err
+	}
+	contents, err := ioutil.ReadAll(buf)
+	if err != nil {
+		return nil, filesAdded, err
+	}
+	return contents, filesAdded, err
+}
+
+// TakeSnapshotOfFiles takes a snapshot of specific files
+// Used for ADD/COPY commands, when we know which files have changed
+func (s *Snapshotter) TakeSnapshotOfFiles(files []string) ([]byte, error) {
+	logrus.Infof("Taking snapshot of files %s", files)
+	s.l.Snapshot()
+	if len(files) == 0 {
+		logrus.Info("No files changed in this command, skipping snapshotting.")
+		return nil, nil
+	}
+	buf := bytes.NewBuffer([]byte{})
+	w := tar.NewWriter(buf)
+	defer w.Close()
+	for _, file := range files {
+		info, err := os.Stat(file)
+		if err != nil {
+			return nil, err
+		}
+		if util.PathInWhitelist(file, s.directory) {
+			logrus.Debugf("Not adding %s to layer, as it is whitelisted", file)
+			continue
+		}
+		// Only add to the tar if we add it to the layeredmap.
+		maybeAdd, err := s.l.MaybeAdd(file)
+		if err != nil {
+			return nil, err
+		}
+		if maybeAdd {
+			util.AddToTar(file, info, w)
+		}
+	}
+	return ioutil.ReadAll(buf)
+}
+
+func (s *Snapshotter) snapShotFS(f io.Writer) (bool, error) {
+	s.l.Snapshot()
+	filesAdded := 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) {
+			logrus.Debugf("Not adding %s to layer, as it's whitelisted", path)
+			return nil
+		}
+
+		// Only add to the tar if we add it to the layeredmap.
+		maybeAdd, err := s.l.MaybeAdd(path)
+		if err != nil {
+			return err
+		}
+		if maybeAdd {
+			filesAdded = true
+			return util.AddToTar(path, info, w)
+		}
+		return nil
+	})
+	return filesAdded, err
+}

pkg/snapshot/snapshot_test.go

@@ -0,0 +1,216 @@
+/*
+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",
+		"bar/bat":       "baz",
+		"workspace/bat": "bat",
+	}
+	if err := testutil.SetupFiles(testDir, newFiles); err != nil {
+		t.Fatalf("Error setting up fs: %s", err)
+	}
+	// Take another snapshot
+	contents, filesAdded, err := snapshotter.TakeSnapshot()
+	if err != nil {
+		t.Fatalf("Error taking snapshot of fs: %s", err)
+	}
+	if !filesAdded {
+		t.Fatal("No files added to snapshot.")
+	}
+	// 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")
+	batPath := filepath.Join(testDir, "bar/bat")
+	snapshotFiles := map[string]string{
+		fooPath: "newbaz1",
+		batPath: "baz",
+	}
+	numFiles := 0
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		numFiles++
+		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))
+		}
+	}
+	if numFiles != 2 {
+		t.Fatalf("Incorrect number of files were added, expected: 2, actual: %v", numFiles)
+	}
+}
+
+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, filesAdded, err := snapshotter.TakeSnapshot()
+	if err != nil {
+		t.Fatalf("Error taking snapshot of fs: %s", err)
+	}
+	if !filesAdded {
+		t.Fatal("No files added to snapshot.")
+	}
+	// 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",
+	}
+	numFiles := 0
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		numFiles++
+		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))
+		}
+	}
+	if numFiles != 1 {
+		t.Fatalf("Incorrect number of files were added, expected: 1, got: %v", numFiles)
+	}
+}
+
+func TestSnapshotFiles(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/file": "bat",
+	}
+	if err := testutil.SetupFiles(testDir, newFiles); err != nil {
+		t.Fatalf("Error setting up fs: %s", err)
+	}
+	filesToSnapshot := []string{
+		filepath.Join(testDir, "foo"),
+		filepath.Join(testDir, "workspace/file"),
+	}
+	contents, err := snapshotter.TakeSnapshotOfFiles(filesToSnapshot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedContents := map[string]string{
+		filepath.Join(testDir, "foo"): "newbaz1",
+	}
+	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
+	reader := bytes.NewReader(contents)
+	tr := tar.NewReader(reader)
+	numFiles := 0
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		numFiles = numFiles + 1
+		if _, isFile := expectedContents[hdr.Name]; !isFile {
+			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
+		}
+		contents, _ := ioutil.ReadAll(tr)
+		if string(contents) != expectedContents[hdr.Name] {
+			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, expectedContents[hdr.Name], string(contents))
+		}
+	}
+	if numFiles != 1 {
+		t.Fatalf("%s was not added.", filepath.Join(testDir, "foo"))
+	}
+}
+
+func TestEmptySnapshot(t *testing.T) {
+	testDir, snapshotter, err := setUpTestDir()
+	defer os.RemoveAll(testDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Take snapshot with no changes
+	_, filesAdded, 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 filesAdded {
+		t.Fatal("Files added even though no changes to file system 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, error) {
+	hasher := func(p string) (string, error) {
+		h := md5.New()
+		fi, err := os.Lstat(p)
+		if err != nil {
+			return "", 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 {
+				return "", err
+			}
+			defer f.Close()
+			if _, err := io.Copy(h, f); err != nil {
+				return "", err
+			}
+		}
+
+		return hex.EncodeToString(h.Sum(nil)), 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)