2 files changed, 90 insertions, 25 deletions
diff --git a/pkg/sentry/device/BUILD b/pkg/sentry/device/BUILD
index 01de708d3..4ccf0674d 100644
--- a/pkg/sentry/device/BUILD
+++ b/pkg/sentry/device/BUILD
@@ -1,7 +1,7 @@
-load("//tools/go_stateify:defs.bzl", "go_library", "go_test")
-
 package(licenses = ["notice"])
 
+load("//tools/go_stateify:defs.bzl", "go_library", "go_test")
+
 go_library(
     name = "device",
     srcs = ["device.go"],
diff --git a/pkg/sentry/device/device.go b/pkg/sentry/device/device.go
index 27e4eb258..ae4fa1d93 100644
--- a/pkg/sentry/device/device.go
+++ b/pkg/sentry/device/device.go
@@ -14,11 +14,6 @@
 
 // Package device defines reserved virtual kernel devices and structures
 // for managing them.
-//
-// Saving and restoring devices is not necessary if the devices are initialized
-// as package global variables. Package initialization happens in a single goroutine
-// and in a deterministic order, so minor device numbers will be assigned in the
-// same order as packages are loaded.
 package device
 
 import (
@@ -30,7 +25,83 @@ import (
 	"gvisor.googlesource.com/gvisor/pkg/abi/linux"
 )
 
+// Registry tracks all simple devices and related state on the system for
+// save/restore.
+//
+// The set of devices across save/restore must remain consistent. That is, no
+// devices may be created or removed on restore relative to the saved
+// system. Practically, this means do not create new devices specifically as
+// part of restore.
+//
+// +stateify savable
+type Registry struct {
+	// lastAnonDeviceMinor is the last minor device number used for an anonymous
+	// device. Must be accessed atomically.
+	lastAnonDeviceMinor uint64
+
+	// mu protects the fields below.
+	mu sync.Mutex `state:"nosave"`
+
+	devices map[ID]*Device
+}
+
+// SimpleDevices is the system-wide simple device registry. This is
+// saved/restored by kernel.Kernel, but defined here to allow access without
+// depending on the kernel package. See kernel.Kernel.deviceRegistry.
+var SimpleDevices = newRegistry()
+
+func newRegistry() *Registry {
+	return &Registry{
+		devices: make(map[ID]*Device),
+	}
+}
+
+// newAnonID assigns a major and minor number to an anonymous device ID.
+func (r *Registry) newAnonID() ID {
+	return ID{
+		// Anon devices always have a major number of 0.
+		Major: 0,
+		// Use the next minor number.
+		Minor: atomic.AddUint64(&r.lastAnonDeviceMinor, 1),
+	}
+}
+
+// newAnonDevice allocates a new anonymous device with a unique minor device
+// number, and registers it with r.
+func (r *Registry) newAnonDevice() *Device {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	d := &Device{
+		ID: r.newAnonID(),
+	}
+	r.devices[d.ID] = d
+	return d
+}
+
+// LoadFrom initializes the internal state of all devices in r from other. The
+// set of devices in both registries must match. Devices may not be created or
+// destroyed across save/restore.
+func (r *Registry) LoadFrom(other *Registry) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	other.mu.Lock()
+	defer other.mu.Unlock()
+	if len(r.devices) != len(other.devices) {
+		panic(fmt.Sprintf("Devices were added or removed when restoring the registry:\nnew:\n%+v\nold:\n%+v", r.devices, other.devices))
+	}
+	for id, otherD := range other.devices {
+		ourD, ok := r.devices[id]
+		if !ok {
+			panic(fmt.Sprintf("Device %+v could not be restored as it wasn't defined in the new registry", otherD))
+		}
+		ourD.loadFrom(otherD)
+	}
+	atomic.StoreUint64(&r.lastAnonDeviceMinor, atomic.LoadUint64(&other.lastAnonDeviceMinor))
+}
+
 // ID identifies a device.
+//
+// +stateify savable
 type ID struct {
 	Major uint64
 	Minor uint64
@@ -41,18 +112,12 @@ func (i *ID) DeviceID() uint64 {
 	return uint64(linux.MakeDeviceID(uint16(i.Major), uint32(i.Minor)))
 }
 
-// nextAnonDeviceMinor is the next minor number for a new anonymous device.
-// Must be accessed atomically.
-var nextAnonDeviceMinor uint64
-
 // NewAnonDevice creates a new anonymous device. Packages that require an anonymous
 // device should initialize the device in a global variable in a file called device.go:
 //
 // var myDevice = device.NewAnonDevice()
 func NewAnonDevice() *Device {
-	return &Device{
-		ID: newAnonID(),
-	}
+	return SimpleDevices.newAnonDevice()
 }
 
 // NewAnonMultiDevice creates a new multi-keyed anonymous device. Packages that require
@@ -62,21 +127,13 @@ func NewAnonDevice() *Device {
 // var myDevice = device.NewAnonMultiDevice()
 func NewAnonMultiDevice() *MultiDevice {
 	return &MultiDevice{
-		ID: newAnonID(),
-	}
-}
-
-// newAnonID assigns a major and minor number to an anonymous device ID.
-func newAnonID() ID {
-	return ID{
-		// Anon devices always have a major number of 0.
-		Major: 0,
-		// Use the next minor number.
-		Minor: atomic.AddUint64(&nextAnonDeviceMinor, 1),
+		ID: SimpleDevices.newAnonID(),
 	}
 }
 
 // Device is a simple virtual kernel device.
+//
+// +stateify savable
 type Device struct {
 	ID
 
@@ -84,6 +141,14 @@ type Device struct {
 	last uint64
 }
 
+// loadFrom initializes d from other. The IDs of both devices must match.
+func (d *Device) loadFrom(other *Device) {
+	if d.ID != other.ID {
+		panic(fmt.Sprintf("Attempting to initialize a device %+v from %+v, but device IDs don't match", d, other))
+	}
+	atomic.StoreUint64(&d.last, atomic.LoadUint64(&other.last))
+}
+
 // NextIno generates a new inode number
 func (d *Device) NextIno() uint64 {
 	return atomic.AddUint64(&d.last, 1)