Do not overwrite handlers in LoadSnapshot

As LoadSnapshot requires only a subset of the default handlers, it was
implemented by assigning the FcInit handlers to a smaller list of
handlers. However, this meant that it would overwrite the FcInit
handlers given by the jailer. This PR fixes this behavior by just
removing the unnecessary handlers, instead of overwriting them entirely.

Signed-off-by: David Son <[email protected]>

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).