Implemented snapshot loading capabilities

Signed-off-by: David Son <davbson@amazon.com>

.buildkite/hooks/pre-exit

@@ -1,4 +1,4 @@
 #!/bin/bash
 set -euo pipefail
 
-find . -user root -print0 | xargs -0 sudo rm -rf '{}'
+sudo chown -hR "${USER}:${USER}" .

.buildkite/pipeline.yml

@@ -78,6 +78,9 @@ steps:
       - ln -s /var/lib/fc-ci/rootfs.ext4 ${FC_TEST_DATA_PATH}/root-drive.img
       # Download Firecracker and its jailer.
       - make deps
+      # Build a rootfs with SSH enabled.
+      - sudo -E FC_TEST_DATA_PATH=${FC_TEST_DATA_PATH} make ${FC_TEST_DATA_PATH}/root-drive-ssh-key
+      - stat ${FC_TEST_DATA_PATH}/root-drive-ssh-key ${FC_TEST_DATA_PATH}/root-drive-with-ssh.img
     agents:
       queue: "${BUILDKITE_AGENT_META_DATA_QUEUE:-default}"
       distro: "${BUILDKITE_AGENT_META_DATA_DISTRO}"
@@ -110,7 +113,7 @@ steps:
 
   - label: ':hammer: root tests'
     commands:
-      - "sudo PATH=$PATH FC_TEST_TAP=fc-root-tap${BUILDKITE_BUILD_NUMBER} FC_TEST_DATA_PATH=${FC_TEST_DATA_PATH} make test EXTRAGOARGS='-v -count=1 -race' DISABLE_ROOT_TESTS="
+      - "sudo -E PATH=$PATH FC_TEST_TAP=fc-root-tap${BUILDKITE_BUILD_NUMBER} FC_TEST_DATA_PATH=${FC_TEST_DATA_PATH} make test EXTRAGOARGS='-v -count=1 -race' DISABLE_ROOT_TESTS="
     agents:
       queue: "${BUILDKITE_AGENT_META_DATA_QUEUE:-default}"
       distro: "${BUILDKITE_AGENT_META_DATA_DISTRO}"

.gitignore

@@ -7,6 +7,7 @@ root-drive.img
 TestPID.img
 build/
 testdata/fc.stamp
+testdata/bin/
 testdata/logs/
 testdata/ltag
 testdata/release-*

Makefile

@@ -20,6 +20,7 @@ FIRECRACKER_DIR=build/firecracker
 FIRECRACKER_TARGET?=x86_64-unknown-linux-musl
 
 FC_TEST_DATA_PATH?=testdata
+FC_TEST_BIN_PATH:=$(FC_TEST_DATA_PATH)/bin
 FIRECRACKER_BIN=$(FC_TEST_DATA_PATH)/firecracker-main
 JAILER_BIN=$(FC_TEST_DATA_PATH)/jailer-main
 
@@ -37,7 +38,17 @@ $(FC_TEST_DATA_PATH)/vmlinux \
 $(FC_TEST_DATA_PATH)/root-drive.img \
 $(FC_TEST_DATA_PATH)/jailer \
 $(FC_TEST_DATA_PATH)/firecracker \
-$(FC_TEST_DATA_PATH)/ltag
+$(FC_TEST_DATA_PATH)/ltag \
+$(FC_TEST_BIN_PATH)/ptp \
+$(FC_TEST_BIN_PATH)/host-local \
+$(FC_TEST_BIN_PATH)/static \
+$(FC_TEST_BIN_PATH)/tc-redirect-tap
+
+# Enable pulling of artifacts from S3 instead of building
+# TODO: https://github.com/firecracker-microvm/firecracker-go-sdk/issues/418
+ifeq ($(GID), 0)
+testdata_objects += $(FC_TEST_DATA_PATH)/root-drive-with-ssh.img $(FC_TEST_DATA_PATH)/root-drive-ssh-key
+endif
 
 testdata_dir = testdata/firecracker.tgz testdata/firecracker_spec-$(firecracker_version).yaml testdata/LICENSE testdata/NOTICE testdata/THIRD-PARTY
 
@@ -82,6 +93,34 @@ $(FC_TEST_DATA_PATH)/fc.stamp:
 $(FC_TEST_DATA_PATH)/root-drive.img:
 	$(curl) -o $@ https://s3.amazonaws.com/spec.ccfc.min/img/hello/fsfiles/hello-rootfs.ext4
 
+$(FC_TEST_DATA_PATH)/root-drive-ssh-key $(FC_TEST_DATA_PATH)/root-drive-with-ssh.img: 
+# Need root to move ssh key to testdata location
+ifeq ($(GID), 0)
+	$(MAKE) $(FIRECRACKER_DIR)
+	$(FIRECRACKER_DIR)/tools/devtool build_rootfs
+	cp $(FIRECRACKER_DIR)/build/rootfs/bionic.rootfs.ext4 $(FC_TEST_DATA_PATH)/root-drive-with-ssh.img
+	cp $(FIRECRACKER_DIR)/build/rootfs/ssh/id_rsa $(FC_TEST_DATA_PATH)/root-drive-ssh-key
+	rm -rf $(FIRECRACKER_DIR)
+else
+	$(error unable to place ssh key without root permissions)
+endif
+
+$(FC_TEST_BIN_PATH)/ptp:
+	GO111MODULE=off GOBIN=$(abspath $(FC_TEST_BIN_PATH)) \
+	go get github.com/containernetworking/plugins/plugins/main/ptp
+
+$(FC_TEST_BIN_PATH)/host-local:
+	GO111MODULE=off GOBIN=$(abspath $(FC_TEST_BIN_PATH)) \
+	go get github.com/containernetworking/plugins/plugins/ipam/host-local
+
+$(FC_TEST_BIN_PATH)/static:
+	GO111MODULE=off GOBIN=$(abspath $(FC_TEST_BIN_PATH)) \
+	go get github.com/containernetworking/plugins/plugins/ipam/static
+
+$(FC_TEST_BIN_PATH)/tc-redirect-tap:
+	GO111MODULE=off GOBIN=$(abspath $(FC_TEST_BIN_PATH)) \
+	go get github.com/awslabs/tc-redirect-tap/cmd/tc-redirect-tap
+
 $(FC_TEST_DATA_PATH)/ltag:
 	GO111MODULE=off GOBIN=$(abspath $(FC_TEST_DATA_PATH)) \
 	go get github.com/kunalkushwaha/ltag

firecracker.go

@@ -281,6 +281,23 @@ func (f *Client) CreateSnapshot(ctx context.Context, snapshotParams *models.Snap
 	return f.client.Operations.CreateSnapshot(params)
 }
 
+// LoadSnapshotOpt is a functional option to be used for the
+// LoadSnapshot API in setting any additional optional fields.
+type LoadSnapshotOpt func(*ops.LoadSnapshotParams)
+
+// LoadSnapshot is a wrapper for the swagger generated client to make
+// calling of the API easier.
+func (f *Client) LoadSnapshot(ctx context.Context, snapshotParams *models.SnapshotLoadParams, opts ...LoadSnapshotOpt) (*ops.LoadSnapshotNoContent, error) {
+	params := ops.NewLoadSnapshotParamsWithContext(ctx)
+	params.SetBody(snapshotParams)
+
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.LoadSnapshot(params)
+}
+
 // CreateSyncActionOpt is a functional option to be used for the
 // CreateSyncAction API in setting any additional optional fields.
 type CreateSyncActionOpt func(*ops.CreateSyncActionParams)

go.mod

@@ -18,5 +18,6 @@ require (
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.8.0
 	github.com/vishvananda/netlink v1.1.1-0.20210330154013-f5de75959ad5
+	golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d
 	golang.org/x/sys v0.0.0-20220204135822-1c1b9b1eba6a
 )

go.sum

@@ -724,6 +724,8 @@ golang.org/x/crypto v0.0.0-20200728195943-123391ffb6de/go.mod h1:LzIPMQfyMNhhGPh
 golang.org/x/crypto v0.0.0-20201002170205-7f63de1d35b0/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201216223049-8b5274cf687f/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
 golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
+golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d h1:sK3txAijHtOK88l68nt020reeT1ZdKLIYetKl95FzVY=
+golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8=
@@ -792,6 +794,7 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v
 golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc=
 golang.org/x/net v0.0.0-20210421230115-4e50805a0758/go.mod h1:72T/g9IO56b78aLF+1Kcs5dz7/ng1VjMUvfKvpfy+jM=
 golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
+golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd h1:O7DYs+zxREGLKzKoMQrtrEacpb0ZVXA5rIwylE2Xchk=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
@@ -886,6 +889,7 @@ golang.org/x/sys v0.0.0-20220204135822-1c1b9b1eba6a h1:ppl5mZgokTT8uPkmYOyEUmPTr
 golang.org/x/sys v0.0.0-20220204135822-1c1b9b1eba6a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

handlers.go

@@ -34,7 +34,8 @@ const (
 	LinkFilesToRootFSHandlerName       = "fcinit.LinkFilesToRootFS"
 	SetupNetworkHandlerName            = "fcinit.SetupNetwork"
 	SetupKernelArgsHandlerName         = "fcinit.SetupKernelArgs"
-	CreateBalloonHandlerName           = "fcint.CreateBalloon"
+	CreateBalloonHandlerName           = "fcinit.CreateBalloon"
+	LoadSnapshotHandlerName            = "fcinit.LoadSnapshot"
 
 	ValidateCfgHandlerName        = "validate.Cfg"
 	ValidateJailerCfgHandlerName  = "validate.JailerCfg"
@@ -280,6 +281,15 @@ func NewCreateBalloonHandler(amountMib int64, deflateOnOom bool, StatsPollingInt
 	}
 }
 
+// LoadSnapshotHandler is a named handler that loads a snapshot
+// from the specified filepath
+var LoadSnapshotHandler = Handler{
+	Name: LoadSnapshotHandlerName,
+	Fn: func(ctx context.Context, m *Machine) error {
+		return m.loadSnapshot(ctx, &m.Cfg.Snapshot)
+	},
+}
+
 var defaultFcInitHandlerList = HandlerList{}.Append(
 	SetupNetworkHandler,
 	SetupKernelArgsHandler,
@@ -294,6 +304,15 @@ var defaultFcInitHandlerList = HandlerList{}.Append(
 	ConfigMmdsHandler,
 )
 
+var loadSnapshotHandlerList = HandlerList{}.Append(
+	SetupNetworkHandler,
+	StartVMMHandler,
+	CreateLogFilesHandler,
+	BootstrapLoggingHandler,
+	LoadSnapshotHandler,
+	AddVsocksHandler,
+)
+
 var defaultValidationHandlerList = HandlerList{}.Append(
 	NetworkConfigValidationHandler,
 )

machine.go

@@ -151,6 +151,13 @@ type Config struct {
 	// It is possible to use a valid IPv4 link-local address (169.254.0.0/16).
 	// If not provided, the default address (169.254.169.254) will be used.
 	MmdsAddress net.IP
+
+	// Configuration for snapshot loading
+	Snapshot SnapshotConfig
+}
+
+func (cfg *Config) hasSnapshot() bool {
+	return cfg.Snapshot.MemFilePath != "" || cfg.Snapshot.SnapshotPath != ""
 }
 
 // Validate will ensure that the required fields are set and that
@@ -381,7 +388,7 @@ func NewMachine(ctx context.Context, cfg Config, opts ...Opt) (*Machine, error)
 // handlers succeed, then this will start the VMM instance.
 // Start may only be called once per Machine.  Subsequent calls will return
 // ErrAlreadyStarted.
-func (m *Machine) Start(ctx context.Context) error {
+func (m *Machine) Start(ctx context.Context, opts ...StartOpt) error {
 	m.logger.Debug("Called Machine.Start()")
 	alreadyStarted := true
 	m.startOnce.Do(func() {
@@ -402,6 +409,10 @@ func (m *Machine) Start(ctx context.Context) error {
 		}
 	}()
 
+	for _, opt := range opts {
+		opt(m)
+	}
+
 	err = m.Handlers.Run(ctx, m)
 	if err != nil {
 		return err
@@ -854,6 +865,10 @@ func (m *Machine) addVsock(ctx context.Context, dev VsockDevice) error {
 }
 
 func (m *Machine) startInstance(ctx context.Context) error {
+	if m.Cfg.hasSnapshot() {
+		return nil
+	}
+
 	action := models.InstanceActionInfoActionTypeInstanceStart
 	info := models.InstanceActionInfo{
 		ActionType: &action,
@@ -1105,6 +1120,23 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 	return nil
 }
 
+// loadSnapshot loads a snapshot of the VM
+func (m *Machine) loadSnapshot(ctx context.Context, snapshot *SnapshotConfig) error {
+	snapshotParams := &models.SnapshotLoadParams{
+		MemFilePath:         &snapshot.MemFilePath,
+		SnapshotPath:        &snapshot.SnapshotPath,
+		EnableDiffSnapshots: snapshot.EnableDiffSnapshots,
+		ResumeVM:            snapshot.ResumeVM,
+	}
+
+	if _, err := m.client.LoadSnapshot(ctx, snapshotParams); err != nil {
+		return fmt.Errorf("failed to load a snapshot for VM: %v", err)
+	}
+
+	m.logger.Debug("snapshot loaded successfully")
+	return nil
+}
+
 // CreateBalloon creates a balloon device if one does not exist
 func (m *Machine) CreateBalloon(ctx context.Context, amountMib int64, deflateOnOom bool, statsPollingIntervals int64, opts ...PutBalloonOpt) error {
 	balloon := models.Balloon{