From 1cd5f14f3afbbb892cd91dfa55aace425ef0aeb9 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Wed, 5 Jun 2024 22:39:51 -0700
Subject: [PATCH] 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 | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 368fec637..daf746b1f 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/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