Merge pull request #630 from sondavidb/fix-snapshot-handler-override

Do not overwrite handlers in LoadSnapshot

Author: sondavidb

handlers.go

@@ -315,24 +315,23 @@ var defaultFcInitHandlerList = HandlerList{}.Append(
 	ConfigMmdsHandler,
 )
 
-var loadSnapshotHandlerList = HandlerList{}.Append(
-	SetupNetworkHandler,
-	StartVMMHandler,
-	CreateLogFilesHandler,
-	BootstrapLoggingHandler,
-	LoadSnapshotHandler,
-	AddVsocksHandler,
-)
+// When the machine starts, these handlers cannot run
+// if we plan to load a snapshot. As these handlers are
+// included in defaultFcInitHandlerList, we must remove them
+// if WithSnapshot() has been specified.
+var loadSnapshotRemoveHandlerList = []Handler{
+	SetupKernelArgsHandler,
+	CreateMachineHandler,
+	CreateBootSourceHandler,
+	AttachDrivesHandler,
+	CreateNetworkInterfacesHandler,
+	ConfigMmdsHandler,
+}
 
 var defaultValidationHandlerList = HandlerList{}.Append(
 	NetworkConfigValidationHandler,
 )
 
-var loadSnapshotValidationHandlerList = HandlerList{}.Append(
-	NetworkConfigValidationHandler,
-	LoadSnapshotConfigValidationHandler,
-)
-
 var defaultHandlers = Handlers{
 	Validation: defaultValidationHandlerList,
 	FcInit:     defaultFcInitHandlerList,

opts.go

@@ -71,11 +71,19 @@ func WithSnapshot(memFilePath, snapshotPath string, opts ...WithSnapshotOpt) Opt
 			opt(&m.Cfg.Snapshot)
 		}
 
-		m.Handlers.Validation = loadSnapshotValidationHandlerList
-		m.Handlers.FcInit = loadSnapshotHandlerList
+		m.Handlers.Validation = m.Handlers.Validation.Append(LoadSnapshotConfigValidationHandler)
+		m.Handlers.FcInit = modifyHandlersForLoadSnapshot(m.Handlers.FcInit)
 	}
 }
 
+func modifyHandlersForLoadSnapshot(l HandlerList) HandlerList {
+	for _, h := range loadSnapshotRemoveHandlerList {
+		l = l.Remove(h.Name)
+	}
+	l = l.Append(LoadSnapshotHandler)
+	return l
+}
+
 // WithMemoryBackend sets the memory backend to the given type, using the given
 // backing file path (a regular file for "File" type, or a UFFD socket path for
 // "Uffd" type).