From ff329286eca6120d71ee71bc14ca7681c504c87d Mon Sep 17 00:00:00 2001
From: Kazuyoshi Kato <katokazu@amazon.com>
Date: Mon, 7 Mar 2022 18:01:32 +0000
Subject: [PATCH] Hide LoadSnapshot from Machine

We should not expose LoadSnapshot, since it is allowed only before boot.

Signed-off-by: Kazuyoshi Kato <katokazu@amazon.com>
---
 machine.go      | 41 +++++++++++++++++++++++++----------------
 machine_test.go |  8 ++++++--
 opts.go         | 11 -----------
 3 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/machine.go b/machine.go
index 14d7445f..a4619e31 100644
--- a/machine.go
+++ b/machine.go
@@ -150,6 +150,12 @@ 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
+
+	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
@@ -722,6 +728,17 @@ func (m *Machine) captureFifoToFileWithChannel(ctx context.Context, logger *log.
 }
 
 func (m *Machine) createMachine(ctx context.Context) error {
+	ss := m.Cfg.Snapshot
+	if ss.SnapshotPath != "" || ss.MemFilePath != "" {
+		_, err := m.client.LoadSnapshot(ctx, &models.SnapshotLoadParams{
+			SnapshotPath:        String(ss.SnapshotPath),
+			MemFilePath:         String(ss.MemFilePath),
+			EnableDiffSnapshots: ss.EnableDiffSnapshots,
+			ResumeVM:            ss.ResumeVM,
+		})
+		return err
+	}
+
 	resp, err := m.client.PutMachineConfiguration(ctx, &m.Cfg.MachineCfg)
 	if err != nil {
 		m.logger.Errorf("PutMachineConfiguration returned %s", resp.Error())
@@ -738,6 +755,10 @@ func (m *Machine) createMachine(ctx context.Context) error {
 }
 
 func (m *Machine) createBootSource(ctx context.Context, imagePath, initrdPath, kernelArgs string) error {
+	if m.Cfg.hasSnapshot() {
+		return nil
+	}
+
 	bsrc := models.BootSource{
 		KernelImagePath: &imagePath,
 		InitrdPath:      initrdPath,
@@ -845,6 +866,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,
@@ -1096,22 +1121,6 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 	return nil
 }
 
-// LoadSnapshot load a snapshot
-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 {
-		m.logger.Errorf("failed to load a snapshot for VM: %v", err)
-		return 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{
diff --git a/machine_test.go b/machine_test.go
index c7f17a0d..c9b3c8a2 100644
--- a/machine_test.go
+++ b/machine_test.go
@@ -1780,7 +1780,11 @@ func TestLoadSnapshot(t *testing.T) {
 
 	// Load a snapshot
 	{
-		cfg := createValidConfig(t, socketPath+".load")
+		cfg := Config{
+			DisableValidation: true,
+			SocketPath:        socketPath + ".load",
+			Snapshot:          SnapshotConfig{SnapshotPath: snapPath, MemFilePath: memPath},
+		}
 		m, err := NewMachine(ctx, cfg, func(m *Machine) {
 			// Rewriting m.cmd partially wouldn't work since Cmd has
 			// some unexported members
@@ -1789,7 +1793,7 @@ func TestLoadSnapshot(t *testing.T) {
 		}, WithLogger(logrus.NewEntry(machineLogger)))
 		require.NoError(t, err)
 
-		err = m.Start(ctx, WithSnapshot(ctx, memPath, snapPath))
+		err = m.Start(ctx)
 		require.NoError(t, err)
 
 		err = m.ResumeVM(ctx)
diff --git a/opts.go b/opts.go
index ba20d640..842772c0 100644
--- a/opts.go
+++ b/opts.go
@@ -14,7 +14,6 @@
 package firecracker
 
 import (
-	"context"
 	"os/exec"
 
 	"github.com/sirupsen/logrus"
@@ -48,13 +47,3 @@ func WithProcessRunner(cmd *exec.Cmd) Opt {
 		machine.cmd = cmd
 	}
 }
-
-// WithSnapshot will allow for the machine to start using a given snapshot.
-func WithSnapshot(ctx context.Context, memFilePath, snapshotPath string) Opt {
-	return func(machine *Machine) {
-		err := machine.LoadSnapshot(ctx, memFilePath, snapshotPath)
-		if err != nil {
-			machine.logger.Errorf("LoadSnapshot failed with %s", err)
-		}
-	}
-}