Implemented snapshot loading capabilities

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

.gitignore

@@ -6,7 +6,7 @@ vmlinux
 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
 
@@ -35,9 +36,22 @@ release_url=https://github.com/firecracker-microvm/firecracker/releases/download
 testdata_objects = \
 $(FC_TEST_DATA_PATH)/vmlinux \
 $(FC_TEST_DATA_PATH)/root-drive.img \
+$(FC_TEST_DATA_PATH)/root-drive-with-ssh.img \
+$(FC_TEST_DATA_PATH)/root-drive-ssh-key \
 $(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 +96,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-with-ssh.img: $(FIRECRACKER_DIR)
+	$(FIRECRACKER_DIR)/tools/devtool build_rootfs
+	cp $(FIRECRACKER_DIR)/build/rootfs/bionic.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)
+	sudo cp $(FIRECRACKER_DIR)/build/rootfs/ssh/id_rsa $@
+else
+	$(warning 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=

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,16 @@ 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 {
+		snapshot := m.Cfg.Snapshot
+		return m.loadSnapshot(ctx, snapshot.MemFilePath, snapshot.SnapshotPath, snapshot.Opts...)
+	},
+}
+
 var defaultFcInitHandlerList = HandlerList{}.Append(
 	SetupNetworkHandler,
 	SetupKernelArgsHandler,
@@ -294,6 +305,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
@@ -826,6 +837,10 @@ func (m *Machine) UpdateGuestNetworkInterfaceRateLimit(ctx context.Context, ifac
 
 // attachDrive attaches a secondary block device
 func (m *Machine) attachDrive(ctx context.Context, dev models.Drive) error {
+	if m.Cfg.hasSnapshot() {
+		return nil
+	}
+
 	hostPath := StringValue(dev.PathOnHost)
 	m.logger.Infof("Attaching drive %s, slot %s, root %t.", hostPath, StringValue(dev.DriveID), BoolValue(dev.IsRootDevice))
 	respNoContent, err := m.client.PutGuestDriveByID(ctx, StringValue(dev.DriveID), &dev)
@@ -854,6 +869,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 +1124,21 @@ 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, memFilePath, snapshotPath string, opts ...LoadSnapshotOpt) error {
+	snapshotParams := &models.SnapshotLoadParams{
+		MemFilePath:  String(memFilePath),
+		SnapshotPath: String(snapshotPath),
+	}
+
+	if _, err := m.client.LoadSnapshot(ctx, snapshotParams, opts...); 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{