From 9f5156a81babb0273b2724cebede17ac7cfab330 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Tue, 21 Apr 2020 08:32:24 -0400
Subject: [PATCH] csi: nil-check allocs for VolumeDenormalize and claim methods
 (#7760)

---
 nomad/state/state_store.go | 10 +++++++---
 nomad/structs/csi.go       |  9 +++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/nomad/state/state_store.go b/nomad/state/state_store.go
index 878f78f4ea2..8e6e644b0b9 100644
--- a/nomad/state/state_store.go
+++ b/nomad/state/state_store.go
@@ -2143,14 +2143,16 @@ func (s *StateStore) CSIVolumeDenormalizePlugins(ws memdb.WatchSet, vol *structs
 	return vol, nil
 }
 
-// csiVolumeDenormalizeAllocs returns a CSIVolume with allocations
+// CSIVolumeDenormalize returns a CSIVolume with allocations
 func (s *StateStore) CSIVolumeDenormalize(ws memdb.WatchSet, vol *structs.CSIVolume) (*structs.CSIVolume, error) {
 	for id := range vol.ReadAllocs {
 		a, err := s.AllocByID(ws, id)
 		if err != nil {
 			return nil, err
 		}
-		vol.ReadAllocs[id] = a
+		if a != nil {
+			vol.ReadAllocs[id] = a
+		}
 	}
 
 	for id := range vol.WriteAllocs {
@@ -2158,7 +2160,9 @@ func (s *StateStore) CSIVolumeDenormalize(ws memdb.WatchSet, vol *structs.CSIVol
 		if err != nil {
 			return nil, err
 		}
-		vol.WriteAllocs[id] = a
+		if a != nil {
+			vol.WriteAllocs[id] = a
+		}
 	}
 
 	return vol, nil
diff --git a/nomad/structs/csi.go b/nomad/structs/csi.go
index c4b650f3a38..9d55a5a844b 100644
--- a/nomad/structs/csi.go
+++ b/nomad/structs/csi.go
@@ -368,6 +368,9 @@ func (v *CSIVolume) Claim(claim CSIVolumeClaimMode, alloc *Allocation) error {
 
 // ClaimRead marks an allocation as using a volume read-only
 func (v *CSIVolume) ClaimRead(alloc *Allocation) error {
+	if alloc == nil {
+		return fmt.Errorf("allocation missing")
+	}
 	if _, ok := v.ReadAllocs[alloc.ID]; ok {
 		return nil
 	}
@@ -385,6 +388,9 @@ func (v *CSIVolume) ClaimRead(alloc *Allocation) error {
 
 // ClaimWrite marks an allocation as using a volume as a writer
 func (v *CSIVolume) ClaimWrite(alloc *Allocation) error {
+	if alloc == nil {
+		return fmt.Errorf("allocation missing")
+	}
 	if _, ok := v.WriteAllocs[alloc.ID]; ok {
 		return nil
 	}
@@ -411,6 +417,9 @@ func (v *CSIVolume) ClaimWrite(alloc *Allocation) error {
 
 // ClaimRelease is called when the allocation has terminated and already stopped using the volume
 func (v *CSIVolume) ClaimRelease(alloc *Allocation) error {
+	if alloc == nil {
+		return fmt.Errorf("allocation missing")
+	}
 	delete(v.ReadAllocs, alloc.ID)
 	delete(v.WriteAllocs, alloc.ID)
 	return nil