Add support for checkpointing additional state in the kernel.

Callers can save additional state into the kernel struct, which will be
checkpointed. The caller package is responsible for consuming this additional
state upon restore/resume.

PiperOrigin-RevId: 640772388

pkg/sentry/kernel/kernel.go

@@ -357,6 +357,10 @@ type Kernel struct {
 	// Mapping: cid -> name.
 	// It's protected by extMu.
 	containerNames map[string]string
+
+	// additionalCheckpointState stores additional state that needs
+	// to be checkpointed. It's protected by extMu.
+	additionalCheckpointState map[any]any
 }
 
 // InitKernelArgs holds arguments to Init.
@@ -1802,6 +1806,28 @@ func (k *Kernel) SetHostMount(mnt *vfs.Mount) {
 	k.hostMount = mnt
 }
 
+// AddStateToCheckpoint adds a key-value pair to be additionally checkpointed.
+func (k *Kernel) AddStateToCheckpoint(key, v any) {
+	k.extMu.Lock()
+	defer k.extMu.Unlock()
+	if k.additionalCheckpointState == nil {
+		k.additionalCheckpointState = make(map[any]any)
+	}
+	k.additionalCheckpointState[key] = v
+}
+
+// PopCheckpointState pops a key-value pair from the additional checkpoint
+// state. If the key doesn't exist, nil is returned.
+func (k *Kernel) PopCheckpointState(key any) any {
+	k.extMu.Lock()
+	defer k.extMu.Unlock()
+	if v, ok := k.additionalCheckpointState[key]; ok {
+		delete(k.additionalCheckpointState, key)
+		return v
+	}
+	return nil
+}
+
 // HostMount returns the hostfs mount.
 func (k *Kernel) HostMount() *vfs.Mount {
 	return k.hostMount